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..2148c6ae 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,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 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()