fix(gateway): a container deploy cancelled mid-create leaves no containers running

src/agent_env/providers/env_providers/env_gateway_provider.py

@@ -935,7 +935,8 @@ async def _deploy_via_containers(
                 self._state_provider.sidecar_specs(instance=state_instance)
                 if self._needs_local_postgres else []
             )
-            # Step 2: MCP server sandboxes (HTTPS), provisioned in parallel.
+            # Step 2: MCP server sandboxes (HTTPS), provisioned in parallel. Each is recorded as it's created, so a
+            # deploy that fails or is cancelled mid-batch closes the ones already made.
             mcp_coros = [
                 self._provision_mcp(deploy, cfg, art, state_instance)
                 for cfg, art in zip(mcp_servers, mcp_server_images)
@@ -954,12 +955,10 @@ async def _deploy_via_containers(
                     continue
                 if i < len(mcp_coros):
                     environment_name, mcp_sb = r
-                    self._container_sandboxes.append(mcp_sb)
                     self._environment_sandboxes[environment_name] = mcp_sb
                     mcp_results.append((environment_name, mcp_sb))
                 else:
                     spec = sidecar_specs[i - len(mcp_coros)]
-                    self._container_sandboxes.append(r)
                     if spec.name == PGWEB_SERVICE_NAME:
                         pgweb_sb = r
                         self._pgweb_sandbox = r
@@ -1106,7 +1105,7 @@ async def _get_container_id(self, sandbox: VmSandbox, compose_service: str) -> s
 
     async def _provision_sidecar(self, deploy: _ContainerDeploy, spec) -> Sandbox:
         logger.info(f"Provisioning {spec.name} from {spec.image}...")
-        return await deploy.sandbox_provider.create_container(
+        sb = await deploy.sandbox_provider.create_container(
             image_name=spec.image,
             port=spec.port,
             env=spec.env,
@@ -1116,6 +1115,8 @@ async def _provision_sidecar(self, deploy: _ContainerDeploy, spec) -> Sandbox:
             attribution=deploy.attribution,
             private_network=True,
         )
+        self._container_sandboxes.append(sb)
+        return sb
 
     @property
     def _needs_local_postgres(self) -> bool:
@@ -1217,6 +1218,7 @@ async def _provision_mcp(
             attribution=deploy.attribution,
             private_network=True,
         )
+        self._container_sandboxes.append(sb)
         return cfg.environment_name, sb
 
     async def _deploy_gateway(

src/agent_env/providers/sandbox_providers/modal_sandbox.py

@@ -5,6 +5,7 @@
 import asyncio
 import base64
 import contextlib
+import functools
 import hashlib
 import logging
 import os
@@ -53,6 +54,36 @@ def _resolve_app_base_name(app_name: Optional[str]) -> str:
 # chunks divisible by 3 so concatenated base64 has padding only at EOF (mid-stream
 # '=' would make `base64 -d` truncate the file).
 _STDIN_CHUNK_BYTES = 8 * 1024 * 1024 - ((8 * 1024 * 1024) % 3)
+# How long a cancelled create waits for its request to Modal to end, a few seconds as a rule; past it the cancel goes
+# on, and the sandbox the request creates is terminated when it ends.
+_CANCELLED_CREATE_WAIT_SECONDS = 30
+# Terminations started after their create's cancel had gone on, kept so they run to their end.
+_LATE_TERMINATIONS: set[asyncio.Task] = set()
+
+
+async def _terminate_created(sb: modal.Sandbox) -> None:
+    """Terminate ``sb``, which a cancelled create made. A failure is logged, naming it, since it runs until its
+    timeout."""
+    try:
+        await sb.terminate.aio()
+    except Exception as e:
+        logger.warning(f"Could not terminate Modal sandbox {sb.object_id}, whose create was cancelled; it runs until "
+                       f"its timeout: {fmt_exc(e)}")
+
+
+def _terminate_when_created(creating: asyncio.Future, *, call_context: str) -> None:
+    """Terminate the sandbox a create that outlasted its cancel's wait made, if it made one. A create cancelled before
+    it ended, as a process that exits cancels it, may have made one all the same, which runs until its timeout, so a
+    warning names the call."""
+    if creating.cancelled():
+        logger.warning(f"A cancelled Modal sandbox create was stopped before Modal answered; a sandbox Modal made for "
+                       f"it runs until its timeout: {call_context}")
+        return
+    if creating.exception() is not None:
+        return
+    task = creating.get_loop().create_task(_terminate_created(creating.result()))
+    _LATE_TERMINATIONS.add(task)
+    task.add_done_callback(_LATE_TERMINATIONS.discard)
 
 
 def _tunnel_url(t) -> str:
@@ -348,15 +379,35 @@ async def create(image: modal.Image) -> modal.Sandbox:
                 *(command or []), i6pn=i6pn, image=image, **create_kwargs
             )
 
-        try:
+        async def create_rebuilding_once() -> modal.Sandbox:
             try:
-                sb = await create(image)
+                return await create(image)
             except modal.exception.NotFoundError:
                 if built_id is None:
                     raise
                 # Modal no longer holds the image this process built from the context, so it's built again, once.
-                image, built_id = await self._context_image(prepared, stale=built_id)
-                sb = await create(image)
+                rebuilt, _ = await self._context_image(prepared, stale=built_id)
+                return await create(rebuilt)
+
+        # Modal may create the sandbox even when this create is cancelled, so the request runs on, and what it creates
+        # is terminated: a cancelled create leaves nothing running. The cancel waits for it a while only, so a request
+        # that hangs can't hold up a chain's deadline or a Ctrl-C.
+        creating = asyncio.ensure_future(create_rebuilding_once())
+        try:
+            sb = await asyncio.shield(creating)
+        except asyncio.CancelledError:
+            created = None
+            try:
+                created = await asyncio.wait_for(asyncio.shield(creating), _CANCELLED_CREATE_WAIT_SECONDS)
+            except (Exception, asyncio.CancelledError):  # it failed, outlasted the wait, or a second cancel came
+                # A create that ended just as the wait gave up runs the callback at once
+                creating.add_done_callback(functools.partial(_terminate_when_created, call_context=call_context))
+                if not creating.done():
+                    logger.warning(f"A cancelled Modal sandbox create is still waiting on Modal; the sandbox it "
+                                   f"creates, if any, is terminated when it ends: {call_context}")
+            if created is not None:
+                await _terminate_created(created)
+            raise
         except Exception as e:
             raise RuntimeError(
                 f"Modal sandbox create failed [{fmt_exc(e)}]: {call_context}"

tst/unit/providers/env_providers/env_gateway_provider_test.py

@@ -574,29 +574,16 @@ async def test_wait_for_tunnel_polls_through_connection_errors(monkeypatch):
     card, seen = await _probe(monkeypatch, httpx.ConnectError("refused"), lambda: httpx.Response(200, json=_CARD))
     assert card == _CARD and len(seen) == 2
 
-async def _run_container_deploy(*, cpu=None, memory_mb=None) -> list[dict]:
-    """Run _deploy_via_containers with all I/O mocked, returning the list of
-    create_container kwargs (one per provisioned component: db, mcp, pgweb,
-    db-mcp, gateway). ECR-style db images are used so the pgweb/db-mcp aux
-    services are provisioned too."""
+@contextlib.contextmanager
+def _slack_in_containers(gp: EnvironmentGatewayProvider, provider: ModalSandboxProvider, create_container):
+    """``gp`` ready to deploy the slack server on ``provider`` in containers, every I/O but ``create_container``
+    mocked. ECR-style db images, so the pgweb and db-mcp sidecars are provisioned too. Yields the deploy's arguments."""
     from agent_env.env.gateway import GatewayMode
 
-    provider = ModalSandboxProvider()
-    gp = EnvironmentGatewayProvider()
     gp._state_provider = LocalPostgresStateProvider()  # normally set by create_gateway
-    create_calls: list[dict] = []
-
-    async def fake_create_container(**kwargs):
-        create_calls.append(kwargs)
-        port = kwargs["port"]
-        image = kwargs["image_name"]
-        if kwargs.get("is_tcp"):
-            return _FakeSandbox(f"db-{port}", port, f"tcp://db.modal.host:{port}")
-        return _FakeSandbox(f"sb-{image}-{port}", port, f"https://{image}.modal.host")
-
     # The servicedb images are in the image store, so the containers run them.
     get_config().set_image_store(EcrImageStore(registry_host="123.dkr.ecr.us-west-2.amazonaws.com"))
-    with patch.object(provider, "create_container", side_effect=fake_create_container):
+    with patch.object(provider, "create_container", side_effect=create_container):
         gp._init_service_db_via_exec = AsyncMock()
         gp._wait_for_tunnel = AsyncMock(return_value=_CARD)
         with patch("agent_env.env.env.Env.get") as env_get, \
@@ -613,22 +600,63 @@ async def fake_create_container(**kwargs):
             cfg.default_service_db_env_id = "db-id"
             get_cfg.return_value = cfg
             env_get.side_effect = lambda env_id, *a, **k: gateway_env if env_id == "gw-id" else db_env
+            yield dict(mcp_servers=[MCPServerConfig(image="mcp-slack", environment_name="slack")],
+                       mcp_server_images=[_image("mcp-slack")], gateway_port=18765, website_configs=None,
+                       gateway_mode=GatewayMode.PERFORMANCE, ttl_seconds=60, disk_size_gb=10)
 
-            await gp._deploy_via_containers(
-                sandbox_provider=provider,
-                mcp_servers=[MCPServerConfig(image="mcp-slack", environment_name="slack")],
-                mcp_server_images=[_image("mcp-slack")],
-                gateway_port=18765,
-                website_configs=None,
-                gateway_mode=GatewayMode.PERFORMANCE,
-                ttl_seconds=60,
-                disk_size_gb=10,
-                cpu=cpu,
-                memory_mb=memory_mb,
-            )
+
+async def _run_container_deploy(*, cpu=None, memory_mb=None) -> list[dict]:
+    """Run _deploy_via_containers with all I/O mocked, returning the list of
+    create_container kwargs (one per provisioned component: db, mcp, pgweb,
+    db-mcp, gateway)."""
+    provider = ModalSandboxProvider()
+    gp = EnvironmentGatewayProvider()
+    create_calls: list[dict] = []
+
+    async def fake_create_container(**kwargs):
+        create_calls.append(kwargs)
+        port = kwargs["port"]
+        image = kwargs["image_name"]
+        if kwargs.get("is_tcp"):
+            return _FakeSandbox(f"db-{port}", port, f"tcp://db.modal.host:{port}")
+        return _FakeSandbox(f"sb-{image}-{port}", port, f"https://{image}.modal.host")
+
+    with _slack_in_containers(gp, provider, fake_create_container) as args:
+        await gp._deploy_via_containers(sandbox_provider=provider, cpu=cpu, memory_mb=memory_mb, **args)
     return create_calls
 
 
+@pytest.mark.asyncio
+async def test_a_deploy_cancelled_while_its_containers_are_created_closes_those_already_made():
+    """Each container is recorded as it's created, so the close a cancelled deploy gets terminates the ones made
+    before the cancel, not only those a finished batch would have recorded."""
+    provider = ModalSandboxProvider()
+    gp = EnvironmentGatewayProvider()
+    created: list[_FakeSandbox] = []
+    made, never = asyncio.Event(), asyncio.Event()
+
+    async def create_container(**kwargs):
+        if kwargs["image_name"].endswith("db-mcp"):
+            await never.wait()  # still being created when the deploy is cancelled
+        name = kwargs["image_name"].rsplit("/", 1)[-1]
+        sandbox = _FakeSandbox(name, kwargs["port"], f"https://{name}.modal.host")
+        created.append(sandbox)
+        if len(created) == 3:  # the servicedb, the slack server and pgweb
+            made.set()
+        return sandbox
+
+    with _slack_in_containers(gp, provider, create_container) as args:
+        deploy = asyncio.create_task(
+            gp._run(provider, "slack", lambda p: gp._deploy_via_containers(sandbox_provider=p, **args)))
+        await asyncio.wait_for(made.wait(), timeout=10)
+        deploy.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await deploy
+
+    assert [sandbox.sandbox_id for sandbox in created] == ["postgres", "mcp-slack", "pgweb"]
+    assert [sandbox.terminated for sandbox in created] == [True, True, True]
+
+
 @pytest.mark.asyncio
 async def test_deploy_level_size_tunes_the_gateway_only():
     """cpu/memory_mb are the gateway's; the per-service sandboxes have their own