fix(sandbox): drain direct release across cancellation (#6348)

* fix(sandbox): drain direct release across cancellation

* test(sandbox): retain direct release across cancellation

* fix(sandbox): drain rollback release across cancellation

* test(sandbox): retain rollback release across cancellation

* test(sandbox): format rollback cancellation regression

backend/packages/harness/deerflow/sandbox/middleware.py

@@ -32,6 +32,7 @@ from deerflow.sandbox.lease import (
 )
 from deerflow.sandbox.overwrite import unwrap_sandbox
 from deerflow.sandbox.sandbox_provider import get_initialized_sandbox_provider
+from deerflow.utils.file_io import await_drained
 
 logger = logging.getLogger(__name__)
 
@@ -223,7 +224,7 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
         if owner_id is not None:
             await get_sandbox_lease_manager(provider).release_async(owner_id)
             return
-        await asyncio.to_thread(provider.release, sandbox_id)
+        await await_drained(asyncio.to_thread(provider.release, sandbox_id))
 
     @override
     def before_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None:

backend/packages/harness/deerflow/sandbox/tools.py

@@ -53,6 +53,7 @@ from deerflow.sandbox.sandbox_provider import SandboxProvider, get_sandbox_provi
 from deerflow.sandbox.search import GrepMatch
 from deerflow.sandbox.security import LOCAL_HOST_BASH_DISABLED_MESSAGE, is_host_bash_allowed
 from deerflow.tools.types import Runtime
+from deerflow.utils.file_io import await_drained
 from deerflow.utils.host_paths import windows_incompatible_segment
 
 logger = logging.getLogger(__name__)
@@ -1606,7 +1607,7 @@ async def _rollback_failed_sandbox_lookup_async(
         if owner_id is not None:
             await get_sandbox_lease_manager(provider).release_async(owner_id)
         else:
-            await asyncio.to_thread(provider.release, sandbox_id)
+            await await_drained(asyncio.to_thread(provider.release, sandbox_id))
     except Exception:
         logger.warning(
             "Failed to roll back sandbox after async post-acquire lookup failure: %s",

backend/tests/test_sandbox_direct_release_cancellation.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import asyncio
+import threading
+
+import pytest
+from langgraph.runtime import Runtime
+
+from deerflow.sandbox.middleware import SandboxMiddleware
+from deerflow.sandbox.sandbox import Sandbox
+from deerflow.sandbox.sandbox_provider import (
+    SandboxProvider,
+    reset_sandbox_provider,
+    set_sandbox_provider,
+)
+
+
+class _BlockingReleaseProvider(SandboxProvider):
+    def __init__(self) -> None:
+        self.release_started = threading.Event()
+        self.allow_release = threading.Event()
+        self.released: list[str] = []
+
+    def acquire(self, thread_id: str | None = None, *, user_id: str | None = None) -> str:
+        raise AssertionError("acquire is not part of this regression")
+
+    def get(self, sandbox_id: str) -> Sandbox | None:
+        return None
+
+    def release(self, sandbox_id: str) -> None:
+        self.release_started.set()
+        assert self.allow_release.wait(timeout=2)
+        self.released.append(sandbox_id)
+
+
+@pytest.mark.anyio
+async def test_aafter_agent_drains_direct_release_across_repeated_cancellation() -> None:
+    provider = _BlockingReleaseProvider()
+    set_sandbox_provider(provider)
+    task: asyncio.Task[dict | None] | None = None
+    try:
+        task = asyncio.create_task(
+            SandboxMiddleware().aafter_agent(
+                {"sandbox": {"sandbox_id": "owned-sandbox"}},
+                Runtime(context={}),
+            )
+        )
+        assert await asyncio.to_thread(provider.release_started.wait, 2)
+
+        task.cancel()
+        await asyncio.sleep(0)
+        task.cancel()
+        await asyncio.sleep(0)
+
+        assert not task.done()
+        provider.allow_release.set()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+    finally:
+        provider.allow_release.set()
+        if task is not None and not task.done():
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+        reset_sandbox_provider()
+
+    assert provider.released == ["owned-sandbox"]

backend/tests/test_sandbox_rollback_release_cancellation.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import asyncio
+import threading
+
+import pytest
+
+from deerflow.sandbox.sandbox import Sandbox
+from deerflow.sandbox.sandbox_provider import SandboxProvider
+from deerflow.sandbox.tools import _rollback_failed_sandbox_lookup_async
+
+
+class _BlockingRollbackProvider(SandboxProvider):
+    def __init__(self) -> None:
+        self.started = threading.Event()
+        self.allow = threading.Event()
+        self.released: list[str] = []
+
+    def acquire(self, thread_id: str | None = None, *, user_id: str | None = None) -> str:
+        raise AssertionError("not used")
+
+    def get(self, sandbox_id: str) -> Sandbox | None:
+        return None
+
+    def release(self, sandbox_id: str) -> None:
+        self.started.set()
+        assert self.allow.wait(timeout=2)
+        self.released.append(sandbox_id)
+
+
+@pytest.mark.anyio
+async def test_async_rollback_drains_direct_release_across_repeated_cancellation() -> None:
+    provider = _BlockingRollbackProvider()
+    task = asyncio.create_task(_rollback_failed_sandbox_lookup_async(provider, "rollback-sandbox", None))
+    try:
+        assert await asyncio.to_thread(provider.started.wait, 2)
+        task.cancel()
+        await asyncio.sleep(0)
+        task.cancel()
+        await asyncio.sleep(0)
+        assert not task.done()
+
+        provider.allow.set()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+    finally:
+        provider.allow.set()
+        if not task.done():
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+
+    assert provider.released == ["rollback-sandbox"]