From a4f39e1b530990c75e1fe791017f2b40f05ed784 Mon Sep 17 00:00:00 2001
From: Ahmed Allam <ahmed39652003@gmail.com>
Date: Thu, 17 Sep 2026 18:45:00 +0000
Subject: [PATCH] fix(runtime): tear the sandbox down when staging is cancelled

CancelledError is not an Exception, so a run cancelled during the extra-file
upload or unpack left a created-but-uncached sandbox running.
---
 strix/runtime/session_manager.py |  2 +-
 tests/test_session_entries.py    | 41 ++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/strix/runtime/session_manager.py b/strix/runtime/session_manager.py
index 6fbdce1e..d908c33d 100644
--- a/strix/runtime/session_manager.py
+++ b/strix/runtime/session_manager.py
@@ -334,7 +334,7 @@ async def create_or_reuse(
         report("Placing workspace files")
         try:
             await stage_extra_files(session, extra_file_archive)
-        except Exception:
+        except BaseException:
             await _discard_session(client, session)
             raise
 
diff --git a/tests/test_session_entries.py b/tests/test_session_entries.py
index 85164b6a..43becfe3 100644
--- a/tests/test_session_entries.py
+++ b/tests/test_session_entries.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import io
 import tarfile
 from dataclasses import dataclass
@@ -479,6 +480,46 @@ async def test_a_failed_unpack_tears_the_session_down(monkeypatch: pytest.Monkey
     assert scan_id not in session_manager._SESSION_CACHE
 
 
+@pytest.mark.asyncio
+async def test_a_cancelled_unpack_tears_the_session_down(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A run cancelled mid-staging must not leave the sandbox running."""
+
+    class _HangingSession(_Session):
+        async def exec(self, *argv: str, timeout: float | None = None) -> _ExecResult:
+            del argv, timeout
+            await asyncio.Event().wait()
+            raise AssertionError("unreachable")
+
+    fake_session = _HangingSession()
+    fake_client = _Client()
+
+    async def _backend(**_kwargs: Any) -> tuple[Any, Any]:
+        return fake_client, fake_session
+
+    scan_id = "unpack-cancelled"
+    backend_name = f"test-{scan_id}"
+    _use_backend(monkeypatch, backend_name, _backend, supports_bind_mounts=True)
+    try:
+        task = asyncio.create_task(
+            session_manager.create_or_reuse(
+                scan_id,
+                image="img",
+                local_sources=[],
+                extra_files=[{"workspace_path": "/workspace/notes.md", "content": b"x"}],
+            )
+        )
+        while not fake_session.writes:
+            await asyncio.sleep(0)
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+    finally:
+        _forget_backend(backend_name)
+
+    assert fake_client.deleted == [fake_session]
+    assert scan_id not in session_manager._SESSION_CACHE
+
+
 def test_only_bind_mount_capable_backends_are_registered_as_such() -> None:
     assert backend_supports_bind_mounts("docker")
     assert not backend_supports_bind_mounts("e2b")