Terminate the VM a gateway deploy created when the deploy fails or is cancelled

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
 

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
 

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
 

src/agent_env/providers/env_providers/env_gateway_provider.py

@@ -607,10 +607,33 @@ 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
         )
+        earlier_sandbox = 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._terminate_new_sandbox(
+                earlier_sandbox=earlier_sandbox, existing_sandbox=gateway.get("existing_sandbox"))
+            await self.close()
+            raise
         return result
 
+    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_sandbox.terminate()
+        except Exception as e:
+            logger.warning(
+                f"Failed to terminate {type(new_sandbox).__name__} {new_sandbox.sandbox_id} after a failed deploy: {e}")
+
     async def _deploy_via_vm(
         self,
         sandbox_provider: SandboxProvider,
@@ -1201,13 +1224,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,
-        )
+        earlier_sandbox = 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._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]]:
         """The containers a container-mode deploy created, in the record's ``sandbox_ids`` shape; ``{}`` on a VM, or only the

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
 

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

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()