From 974d7d8b3151f4fd026d8d217bacf394ebac2579 Mon Sep 17 00:00:00 2001
From: Ahmed Allam <ahmed39652003@gmail.com>
Date: Tue, 13 Oct 2026 15:36:25 +0000
Subject: [PATCH] fix(runtime): keep the default path blocking on the proxy
 endpoint and discard a sandbox created after a cancelled deferred bring-up

---
 strix/runtime/session_manager.py | 48 ++++++++++++++++++++++++++------
 tests/test_deferred_session.py   | 20 ++++++++++++-
 2 files changed, 58 insertions(+), 10 deletions(-)

diff --git a/strix/runtime/session_manager.py b/strix/runtime/session_manager.py
index cdb9a035..cfa5bc4c 100644
--- a/strix/runtime/session_manager.py
+++ b/strix/runtime/session_manager.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import io
 import logging
 import os
@@ -23,6 +24,8 @@ from strix.runtime.deferred_session import DeferredSandboxClient, DeferredSandbo
 
 
 if TYPE_CHECKING:
+    from collections.abc import Awaitable
+
     from agents.sandbox.session import BaseSandboxSession
 
     from strix.runtime.status import StatusSink
@@ -328,11 +331,13 @@ async def create_or_reuse(
 
     async def _bring_up() -> tuple[Any, Any]:
         report("Starting sandbox container")
-        client, session = await backend(
-            image=image,
-            manifest=manifest,
-            exposed_ports=(_CONTAINER_CAIDO_PORT,),
-            bind_mounts=bind_mounts,
+        client, session = await _create_or_discard(
+            backend(
+                image=image,
+                manifest=manifest,
+                exposed_ports=(_CONTAINER_CAIDO_PORT,),
+                bind_mounts=bind_mounts,
+            )
         )
         if extra_file_archive is not None:
             report("Placing workspace files")
@@ -343,13 +348,16 @@ async def create_or_reuse(
                 raise
         return client, session
 
-    async def _caido(session: Any) -> Any:
+    async def _caido_url(session: Any) -> str:
         report("Setting up the proxy")
         caido_endpoint = await session.resolve_exposed_port(_CONTAINER_CAIDO_PORT)
         scheme = "https" if caido_endpoint.tls else "http"
         host_caido_url = f"{scheme}://{caido_endpoint.host}:{caido_endpoint.port}"
         logger.debug("Caido host endpoint resolved: %s", host_caido_url)
-        return await bootstrap_caido(
+        return host_caido_url
+
+    def _caido(session: Any, host_caido_url: str) -> Any:
+        return bootstrap_caido(
             session,
             host_url=host_caido_url,
             container_url=container_caido_url,
@@ -364,7 +372,8 @@ async def create_or_reuse(
         )
 
         async def _caido_when_up() -> Any:
-            return await _caido(await deferred.resolve())
+            session = await deferred.resolve()
+            return await _caido(session, await _caido_url(session))
 
         bundle = {
             "client": DeferredSandboxClient(deferred),
@@ -378,8 +387,9 @@ async def create_or_reuse(
         return bundle
 
     client, session = await _bring_up()
+    host_caido_url = await _caido_url(session)
     caido_client = CaidoBootstrapHandle(
-        asyncio.create_task(_caido(session), name=f"caido-bootstrap-{scan_id}")
+        asyncio.create_task(_caido(session, host_caido_url), name=f"caido-bootstrap-{scan_id}")
     )
 
     bundle = {
@@ -392,6 +402,26 @@ async def create_or_reuse(
     return bundle
 
 
+_DISCARD_TASKS: set[asyncio.Task[None]] = set()
+
+
+async def _create_or_discard(create: Awaitable[tuple[Any, Any]]) -> tuple[Any, Any]:
+    pending = asyncio.ensure_future(create)
+    try:
+        return await asyncio.shield(pending)
+    except asyncio.CancelledError:
+        discard = asyncio.create_task(_discard_when_created(pending))
+        _DISCARD_TASKS.add(discard)
+        discard.add_done_callback(_DISCARD_TASKS.discard)
+        raise
+
+
+async def _discard_when_created(create: asyncio.Future[tuple[Any, Any]]) -> None:
+    with contextlib.suppress(BaseException):
+        client, session = await create
+        await _discard_session(client, session)
+
+
 async def _discard_session(client: Any, session: Any) -> None:
     """Best-effort teardown of a session that never made it into the cache."""
     try:
diff --git a/tests/test_deferred_session.py b/tests/test_deferred_session.py
index 12d16725..aac5d226 100644
--- a/tests/test_deferred_session.py
+++ b/tests/test_deferred_session.py
@@ -129,6 +129,11 @@ def backend(monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest) ->
     _BIND_MOUNT_BACKENDS.discard(name)
 
 
+async def _settle() -> None:
+    for _ in range(10):
+        await asyncio.sleep(0)
+
+
 async def _create(scan_id: str, **kwargs: Any) -> dict[str, Any]:
     return await session_manager.create_or_reuse(
         scan_id, image="img", local_sources=[], defer=True, **kwargs
@@ -144,7 +149,7 @@ async def test_deferred_bundle_returns_before_the_backend_finishes(backend: _Bac
         assert isinstance(session, DeferredSandboxSession)
         assert isinstance(bundle["client"], DeferredSandboxClient)
         assert session.pending
-        await asyncio.sleep(0)
+        await _settle()
         assert backend.calls == 1
         assert await session.running()
         assert session.state.manifest.environment.value["PYTHONUNBUFFERED"] == "1"
@@ -205,6 +210,19 @@ async def test_cleanup_while_pending_cancels_the_bring_up(backend: _Backend) ->
     assert session.bring_up.cancelled()
     assert backend.client.deleted == []
     assert session_manager._SESSION_CACHE.get(scan_id) is None
+    backend.gate.set()
+    await _settle()
+    assert backend.client.deleted == [backend.session]
+
+
+@pytest.mark.asyncio
+async def test_cleanup_while_pending_ignores_a_failed_late_bring_up(backend: _Backend) -> None:
+    scan_id = "deferred-cancel-fail"
+    await _create(scan_id)
+    await asyncio.wait_for(session_manager.cleanup(scan_id), timeout=1)
+    backend.gate.set()
+    await _settle()
+    assert backend.client.deleted == []
 
 
 @pytest.mark.asyncio