From a0e19c97a76067c9c04da9e551aa5d441660759a Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 08:39:29 -0700 Subject: [PATCH 1/3] Terminate the VM a gateway deploy created when the deploy fails or is cancelled A VM-path gateway deploy that failed or was cancelled after create_vm left the VM running, or on the local provider its compose stack: no record named it yet, and close() ends only container sandboxes. Each attempt now terminates the VM it created when it ends that way, as does a cancellation during the tools probe, but never the caller's existing sandbox or one an earlier deploy made. A chained provider's failed members no longer each leave theirs. A cancelled deploy is closed like a failed one, and an env's close() lets a cancellation through, keeping its sandbox for the next close(). Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/env/envs/mcp_server.py | 4 +- src/agent_env/env/envs/multi_env.py | 2 +- src/agent_env/env/envs/website.py | 4 +- .../env_providers/env_gateway_provider.py | 38 +++++++++--- .../providers/env_providers/env_provider.py | 5 +- tst/unit/env/envs/test_env_close_teardown.py | 24 ++++++++ .../env_gateway_provider_test.py | 61 +++++++++++++++++++ 7 files changed, 123 insertions(+), 15 deletions(-) diff --git a/src/agent_env/env/envs/mcp_server.py b/src/agent_env/env/envs/mcp_server.py index 238db5aa..46faa63a 100644 --- a/src/agent_env/env/envs/mcp_server.py +++ b/src/agent_env/env/envs/mcp_server.py @@ -129,12 +129,12 @@ async def close(self) -> None: if self._env_provider is not None: try: await self._env_provider.close() - except BaseException as e: + except Exception as e: logger.warning(f"Failed to close env provider: {e}") if self._sandbox is not None: try: await self._sandbox.terminate() - except BaseException as e: + except Exception as e: # a cancellation propagates, and keeps the sandbox for the next close() logger.warning(f"Failed to terminate sandbox {self._sandbox.sandbox_id}: {e}") self._sandbox = None diff --git a/src/agent_env/env/envs/multi_env.py b/src/agent_env/env/envs/multi_env.py index 040db384..0877a233 100644 --- a/src/agent_env/env/envs/multi_env.py +++ b/src/agent_env/env/envs/multi_env.py @@ -781,7 +781,7 @@ async def close(self) -> None: if self._sandbox is not None: try: await self._sandbox.terminate() - except BaseException as e: + except Exception as e: # a cancellation propagates, and keeps the sandbox for the next close() logger.warning(f"Failed to terminate sandbox {self._sandbox.sandbox_id}: {e}") self._sandbox = None diff --git a/src/agent_env/env/envs/website.py b/src/agent_env/env/envs/website.py index 5d693bb4..da8e6637 100644 --- a/src/agent_env/env/envs/website.py +++ b/src/agent_env/env/envs/website.py @@ -210,12 +210,12 @@ async def close(self) -> None: if self._env_provider is not None: try: await self._env_provider.close() - except BaseException as e: + except Exception as e: logger.warning(f"Failed to close env provider: {e}") if self._sandbox is not None: try: await self._sandbox.terminate() - except BaseException as e: + except Exception as e: # a cancellation propagates, and keeps the sandbox for the next close() logger.warning(f"Failed to terminate sandbox {self._sandbox.sandbox_id}: {e}") self._sandbox = None diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index a842203c..d429e023 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -607,10 +607,27 @@ async def _launch(self, sandbox_provider: SandboxProvider, env_id: str | None, * self._deploy_gateway, **gateway, mcp_server_name=mcp_server_name or random_mcp_server_name(), # drawn once, so every chained attempt presents the same name ) + before = self._sandbox result = await self._run(sandbox_provider, env_id, attempt) - await _probe_tools(env_id, result) + try: + await _probe_tools(env_id, result) + except BaseException: # a cancellation: the probe itself never raises + await self._discard_vm(before, gateway.get("existing_sandbox")) + await self.close() + raise return result + async def _discard_vm(self, before: Sandbox | None, kept: Sandbox | None) -> None: + """Terminate the VM a deploy that failed or was cancelled created, which no record names yet, so nothing else + would. Not one an earlier deploy made (``before``), the caller's (``kept``), or a container, which close() ends.""" + created, self._sandbox = self._sandbox, before + if created is None or created is before or created is kept or created in self._container_sandboxes: + return + try: + await created.terminate() + except Exception as e: + logger.warning(f"Failed to terminate {type(created).__name__} {created.sandbox_id} after a failed deploy: {e}") + async def _deploy_via_vm( self, sandbox_provider: SandboxProvider, @@ -1201,13 +1218,18 @@ async def _deploy_gateway( ttl_seconds=ttl_seconds, disk_size_gb=disk_size_gb, cpu=cpu, memory_mb=memory_mb, mcp_server_name=mcp_server_name, **deploy_kwargs, ) - return await self._deploy_via_vm( - sandbox_provider, mcp_servers, mcp_server_images, - gateway_port=gateway_port, website_configs=website_configs, website_images=website_images, - gateway_mode=gateway_mode, ttl_seconds=ttl_seconds, disk_size_gb=disk_size_gb, - cpu=cpu, memory_mb=memory_mb, existing_sandbox=existing_sandbox, sidecars=sidecars, - mcp_server_name=mcp_server_name, **deploy_kwargs, - ) + before = self._sandbox + try: + return await self._deploy_via_vm( + sandbox_provider, mcp_servers, mcp_server_images, + gateway_port=gateway_port, website_configs=website_configs, website_images=website_images, + gateway_mode=gateway_mode, ttl_seconds=ttl_seconds, disk_size_gb=disk_size_gb, + cpu=cpu, memory_mb=memory_mb, existing_sandbox=existing_sandbox, sidecars=sidecars, + mcp_server_name=mcp_server_name, **deploy_kwargs, + ) + except BaseException: + await self._discard_vm(before, existing_sandbox) + raise def _sandbox_ids(self, always_name_gateway: bool = False) -> dict[str, str | dict[str, str]]: """The containers a container-mode deploy created, in the record's ``sandbox_ids`` shape; ``{}`` on a VM, or only the diff --git a/src/agent_env/providers/env_providers/env_provider.py b/src/agent_env/providers/env_providers/env_provider.py index dba0479e..4846ec01 100644 --- a/src/agent_env/providers/env_providers/env_provider.py +++ b/src/agent_env/providers/env_providers/env_provider.py @@ -147,7 +147,8 @@ def _sandbox_providers(self, sandbox_provider: SandboxProvider, env_id: str | No return [sandbox_provider] async def _run(self, sandbox_provider: SandboxProvider, env_id: str | None, attempt: Callable[[SandboxProvider], Awaitable[Any]]) -> Any: - """Run ``attempt`` on the sandbox provider, or on each chained member until one succeeds; a failed deploy is closed.""" + """Run ``attempt`` on the sandbox provider, or on each chained member until one succeeds; a deploy that fails or is + cancelled is closed.""" try: sandbox_providers = self._sandbox_providers(sandbox_provider, env_id) if len(sandbox_providers) == 1: @@ -161,7 +162,7 @@ async def _run(self, sandbox_provider: SandboxProvider, env_id: str | None, atte await self.close() detail = "; ".join(f"{n}: {e!r}" for n, e in errors) raise RuntimeError(f"All {len(sandbox_providers)} chained providers failed to deploy: {detail}") - except Exception: + except BaseException: await self.close() raise diff --git a/tst/unit/env/envs/test_env_close_teardown.py b/tst/unit/env/envs/test_env_close_teardown.py index 6d17f216..f1477b56 100644 --- a/tst/unit/env/envs/test_env_close_teardown.py +++ b/tst/unit/env/envs/test_env_close_teardown.py @@ -1,10 +1,12 @@ """Unit tests for env close() routing teardown through the gateway provider.""" +import asyncio from unittest.mock import AsyncMock import pytest from agent_env.env.envs.mcp_server import MCPServerEnv +from agent_env.env.envs.multi_env import MultiEnv from agent_env.env.envs.website import WebsiteEnv @@ -74,3 +76,25 @@ async def test_an_env_that_never_deployed_closes_its_own_provider(make_env): env = make_env() await env.close() # the provider built in __init__, not one a test set assert env._sandbox is None + + +def _multi_env() -> MultiEnv: + return MultiEnv(id="crm-suite", version=1, mcp_server_envs=[_mcp_env()]) + + +@pytest.mark.parametrize("make_env", [_mcp_env, _website_env, _multi_env]) +@pytest.mark.asyncio +async def test_a_close_cancelled_while_it_terminates_keeps_the_sandbox_for_the_next(make_env): + env = make_env() + env._env_provider = AsyncMock() + sandbox = AsyncMock() + sandbox.sandbox_id = "sb-1" + sandbox.terminate.side_effect = [asyncio.CancelledError(), None] + env._sandbox = sandbox + + with pytest.raises(asyncio.CancelledError): + await env.close() + assert env._sandbox is sandbox + + await env.close() + assert sandbox.terminate.await_count == 2 and env._sandbox is None diff --git a/tst/unit/providers/env_providers/env_gateway_provider_test.py b/tst/unit/providers/env_providers/env_gateway_provider_test.py index 55804337..9845bac2 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -1419,6 +1419,67 @@ async def run(*args, **kwargs): assert messages[2].startswith("env_deploy_tools_probe env_id=e1 status=ok ") +def _vm(sandbox_id: str) -> _FakeSandbox: + vm = _FakeSandbox(sandbox_id, 8080, f"http://{sandbox_id}") + vm.mode = "vm" + return vm + + +def _creates(gp: EnvironmentGatewayProvider, vm: _FakeSandbox, outcome): + """A VM deploy path that creates ``vm``, then raises ``outcome`` or returns it.""" + async def run(*args, existing_sandbox=None, **kwargs): + gp._sandbox = existing_sandbox or vm + if isinstance(outcome, BaseException): + raise outcome + return outcome + return run + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ended_by", [RuntimeError("compose up failed"), asyncio.CancelledError()], ids=["fails", "cancelled"]) +async def test_a_vm_deploy_that_ends_after_creating_its_vm_terminates_it(ended_by): + gp, vm = EnvironmentGatewayProvider(), _vm("vm-1") + with patch.object(gp, "_deploy_via_vm", side_effect=_creates(gp, vm, ended_by)): + with pytest.raises(type(ended_by)): + await gp.create_gateway(sandbox_provider=MagicMock(), mcp_servers=[], mcp_server_images=[], env_id="e1") + assert vm.terminated and gp.sandbox is None + + +@pytest.mark.asyncio +async def test_a_failed_deploy_leaves_the_callers_vm_alone(): + gp, callers = EnvironmentGatewayProvider(), _vm("callers-vm") + with patch.object(gp, "_deploy_via_vm", side_effect=_creates(gp, _vm("unused"), RuntimeError("boom"))): + with pytest.raises(RuntimeError): + await gp.create_gateway(sandbox_provider=MagicMock(), mcp_servers=[], mcp_server_images=[], env_id="e1", + existing_sandbox=callers) + assert not callers.terminated + + +@pytest.mark.asyncio +async def test_each_failed_chained_member_terminates_its_own_vm_and_the_winner_keeps_its(monkeypatch): + monkeypatch.setattr(env_gateway_provider, "_tool_names", AsyncMock(return_value=set())) + gp, first, second, deployed = EnvironmentGatewayProvider(), _vm("vm-1"), _vm("vm-2"), _probed() + paths = iter([_creates(gp, first, RuntimeError("first")), _creates(gp, second, deployed)]) + + async def each_member(*args, **kwargs): + return await next(paths)(*args, **kwargs) + + with patch.object(gp, "_deploy_via_vm", side_effect=each_member): + chain = ChainedSandboxProvider([MagicMock(), MagicMock()]) + assert await gp.create_gateway(sandbox_provider=chain, mcp_servers=[], mcp_server_images=[], env_id="e1") is deployed + assert first.terminated and not second.terminated and gp.sandbox is second + + +@pytest.mark.asyncio +async def test_a_deploy_cancelled_while_its_tools_are_probed_terminates_its_vm(monkeypatch): + monkeypatch.setattr(env_gateway_provider, "_tool_names", AsyncMock(side_effect=asyncio.CancelledError())) + gp, vm = EnvironmentGatewayProvider(), _vm("vm-1") + with patch.object(gp, "_deploy_via_vm", side_effect=_creates(gp, vm, _probed())): + with pytest.raises(asyncio.CancelledError): + await gp.create_gateway(sandbox_provider=MagicMock(), mcp_servers=[], mcp_server_images=[], env_id="e1") + assert vm.terminated and gp.sandbox is None + + @pytest.mark.asyncio async def test_a_chained_gateway_that_fails_everywhere_names_every_member(): gp = EnvironmentGatewayProvider() From fb36338fb8ad7c46cc576a83d53a4898a53e2787 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 08:49:51 -0700 Subject: [PATCH 2/3] Name the cleanup for what it ends: _terminate_new_vm, given the earlier and the caller's sandboxes by name Co-Authored-By: Claude Opus 5.5 (1M context) --- .../env_providers/env_gateway_provider.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index d429e023..169203e4 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -607,26 +607,30 @@ async def _launch(self, sandbox_provider: SandboxProvider, env_id: str | None, * self._deploy_gateway, **gateway, mcp_server_name=mcp_server_name or random_mcp_server_name(), # drawn once, so every chained attempt presents the same name ) - before = self._sandbox + earlier_sandbox = self._sandbox result = await self._run(sandbox_provider, env_id, attempt) try: await _probe_tools(env_id, result) except BaseException: # a cancellation: the probe itself never raises - await self._discard_vm(before, gateway.get("existing_sandbox")) + await self._terminate_new_vm( + earlier_sandbox=earlier_sandbox, existing_sandbox=gateway.get("existing_sandbox")) await self.close() raise return result - async def _discard_vm(self, before: Sandbox | None, kept: Sandbox | None) -> None: - """Terminate the VM a deploy that failed or was cancelled created, which no record names yet, so nothing else - would. Not one an earlier deploy made (``before``), the caller's (``kept``), or a container, which close() ends.""" - created, self._sandbox = self._sandbox, before - if created is None or created is before or created is kept or created in self._container_sandboxes: + async def _terminate_new_vm(self, *, earlier_sandbox: Sandbox | None, existing_sandbox: Sandbox | None) -> None: + """Terminate the VM a deploy that failed or was cancelled created: no record names it yet, so nothing else + would. Leaves alone an earlier deploy's sandbox, the caller's existing_sandbox, and containers, which close() + ends.""" + new_vm, self._sandbox = self._sandbox, earlier_sandbox + if (new_vm is None or new_vm is earlier_sandbox or new_vm is existing_sandbox + or new_vm in self._container_sandboxes): return try: - await created.terminate() + await new_vm.terminate() except Exception as e: - logger.warning(f"Failed to terminate {type(created).__name__} {created.sandbox_id} after a failed deploy: {e}") + logger.warning( + f"Failed to terminate {type(new_vm).__name__} {new_vm.sandbox_id} after a failed deploy: {e}") async def _deploy_via_vm( self, @@ -1218,7 +1222,7 @@ async def _deploy_gateway( ttl_seconds=ttl_seconds, disk_size_gb=disk_size_gb, cpu=cpu, memory_mb=memory_mb, mcp_server_name=mcp_server_name, **deploy_kwargs, ) - before = self._sandbox + earlier_sandbox = self._sandbox try: return await self._deploy_via_vm( sandbox_provider, mcp_servers, mcp_server_images, @@ -1228,7 +1232,7 @@ async def _deploy_gateway( mcp_server_name=mcp_server_name, **deploy_kwargs, ) except BaseException: - await self._discard_vm(before, existing_sandbox) + await self._terminate_new_vm(earlier_sandbox=earlier_sandbox, existing_sandbox=existing_sandbox) raise def _sandbox_ids(self, always_name_gateway: bool = False) -> dict[str, str | dict[str, str]]: From 12befafd8661f03ec4d9d7835f89489b1065590d Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 09:52:31 -0700 Subject: [PATCH 3/3] Call it _terminate_new_sandbox: what kind of sandbox it ends is the provider's concern Co-Authored-By: Claude Opus 5.5 (1M context) --- .../env_providers/env_gateway_provider.py | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index 169203e4..2148c6ae 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -612,25 +612,27 @@ async def _launch(self, sandbox_provider: SandboxProvider, env_id: str | None, * try: await _probe_tools(env_id, result) except BaseException: # a cancellation: the probe itself never raises - await self._terminate_new_vm( + await self._terminate_new_sandbox( earlier_sandbox=earlier_sandbox, existing_sandbox=gateway.get("existing_sandbox")) await self.close() raise return result - async def _terminate_new_vm(self, *, earlier_sandbox: Sandbox | None, existing_sandbox: Sandbox | None) -> None: - """Terminate the VM a deploy that failed or was cancelled created: no record names it yet, so nothing else - would. Leaves alone an earlier deploy's sandbox, the caller's existing_sandbox, and containers, which close() - ends.""" - new_vm, self._sandbox = self._sandbox, earlier_sandbox - if (new_vm is None or new_vm is earlier_sandbox or new_vm is existing_sandbox - or new_vm in self._container_sandboxes): + async def _terminate_new_sandbox( + self, *, earlier_sandbox: Sandbox | None, existing_sandbox: Sandbox | None, + ) -> None: + """Terminate the sandbox a deploy that failed or was cancelled created: no record names it yet, so nothing else + would. Leaves alone an earlier deploy's sandbox, the caller's existing_sandbox, and the container sandboxes + close() ends.""" + new_sandbox, self._sandbox = self._sandbox, earlier_sandbox + if (new_sandbox is None or new_sandbox is earlier_sandbox or new_sandbox is existing_sandbox + or new_sandbox in self._container_sandboxes): return try: - await new_vm.terminate() + await new_sandbox.terminate() except Exception as e: logger.warning( - f"Failed to terminate {type(new_vm).__name__} {new_vm.sandbox_id} after a failed deploy: {e}") + f"Failed to terminate {type(new_sandbox).__name__} {new_sandbox.sandbox_id} after a failed deploy: {e}") async def _deploy_via_vm( self, @@ -1232,7 +1234,7 @@ async def _deploy_gateway( mcp_server_name=mcp_server_name, **deploy_kwargs, ) except BaseException: - await self._terminate_new_vm(earlier_sandbox=earlier_sandbox, existing_sandbox=existing_sandbox) + await self._terminate_new_sandbox(earlier_sandbox=earlier_sandbox, existing_sandbox=existing_sandbox) raise def _sandbox_ids(self, always_name_gateway: bool = False) -> dict[str, str | dict[str, str]]: