From 23301318fff809abe64ef54466d95c091f080491 Mon Sep 17 00:00:00 2001 From: NanPan <111261006+poijygfdyy@users.noreply.github.com> Date: Tue, 6 Oct 2026 11:22:46 +0800 Subject: [PATCH] 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 --- .../harness/deerflow/sandbox/middleware.py | 3 +- .../harness/deerflow/sandbox/tools.py | 3 +- ...est_sandbox_direct_release_cancellation.py | 66 +++++++++++++++++++ ...t_sandbox_rollback_release_cancellation.py | 52 +++++++++++++++ 4 files changed, 122 insertions(+), 2 deletions(-) create mode 100644 backend/tests/test_sandbox_direct_release_cancellation.py create mode 100644 backend/tests/test_sandbox_rollback_release_cancellation.py diff --git a/backend/packages/harness/deerflow/sandbox/middleware.py b/backend/packages/harness/deerflow/sandbox/middleware.py index 924f801a0a..b574d97a84 100644 --- a/backend/packages/harness/deerflow/sandbox/middleware.py +++ b/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: diff --git a/backend/packages/harness/deerflow/sandbox/tools.py b/backend/packages/harness/deerflow/sandbox/tools.py index 90a57064df..6655be443e 100644 --- a/backend/packages/harness/deerflow/sandbox/tools.py +++ b/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", diff --git a/backend/tests/test_sandbox_direct_release_cancellation.py b/backend/tests/test_sandbox_direct_release_cancellation.py new file mode 100644 index 0000000000..9e4a320702 --- /dev/null +++ b/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"] diff --git a/backend/tests/test_sandbox_rollback_release_cancellation.py b/backend/tests/test_sandbox_rollback_release_cancellation.py new file mode 100644 index 0000000000..0b1e00ff6f --- /dev/null +++ b/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"]