fix(sandbox): retain AIO acquire ownership across cancellation (#5888)

* fix(sandbox): retain AIO acquire ownership across cancellation

* test(sandbox): pin cancelled AIO acquire ownership

* docs(sandbox): document acquire worker ownership

* fix(sandbox): drain only started AIO acquire workers

* test(sandbox): cover queued and started acquire workers

* docs(sandbox): scope AIO acquire ownership guidance

* style(sandbox): separate acquire worker helper

* test(sandbox): queue reclaim deterministically

* style(sandbox): format acquire worker helper

* docs(sandbox): trim acquire guidance below chain limit

* fix(sandbox): separate AIO holder work from lock waiters

* test(sandbox): pin AIO holder executor separation

* docs(sandbox): separate AIO waiter and holder workers

* docs(sandbox): trim AIO worker ownership guidance

* test(sandbox): align acquire worker executor fixtures

* test(sandbox): fix reconciliation executor fixture and formatting

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>

backend/packages/harness/deerflow/community/aio_sandbox/aio_sandbox_provider.py

@@ -13,6 +13,7 @@ The provider itself handles:
 import asyncio
 import atexit
 import contextlib
+import contextvars
 import errno
 import hashlib
 import logging
@@ -23,6 +24,8 @@ import threading
 import time
 import uuid
 from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
 from typing import Any
 
 try:
@@ -117,6 +120,42 @@ class SandboxIdentityCollisionError(RuntimeError):
         self.sandbox_id = sandbox_id
 
 
+async def _run_started_acquire_worker[T](
+    executor,
+    func: Callable[..., T],
+    /,
+    *args: Any,
+    **kwargs: Any,
+) -> T:
+    """Cancel queued acquire work, but drain a worker once it has started."""
+    loop = asyncio.get_running_loop()
+    done = asyncio.Event()
+    context = contextvars.copy_context()
+    call = partial(func, *args, **kwargs)
+    worker = executor.submit(context.run, call)
+    worker.add_done_callback(lambda _future: loop.call_soon_threadsafe(done.set))
+    wrapped = asyncio.wrap_future(worker, loop=loop)
+    try:
+        return await wrapped
+    except asyncio.CancelledError as cancellation:
+        # Awaiting wrapped already asks the concurrent future to cancel. If it
+        # was still queued, preserve the old to_thread behavior: it never runs.
+        if worker.cancelled() or worker.cancel():
+            raise
+        # Once running, the worker cannot be stopped safely. Keep same-scope
+        # serializer ownership until it settles, absorbing repeated cancellation.
+        while not worker.done():
+            try:
+                await asyncio.shield(done.wait())
+            except asyncio.CancelledError:
+                continue
+        try:
+            worker.result()
+        except Exception:
+            logger.warning("Cancelled AIO acquire worker failed while draining", exc_info=True)
+        raise cancellation
+
+
 def _lock_file_exclusive(lock_file) -> None:
     if fcntl is not None:
         fcntl.flock(lock_file, fcntl.LOCK_EX)
@@ -205,6 +244,11 @@ class AioSandboxProvider(WarmPoolLifecycleMixin[SandboxInfo], SandboxProvider):
         self._sandbox_infos: dict[str, SandboxInfo] = {}  # sandbox_id -> SandboxInfo (for destroy)
         self._thread_sandboxes: dict[tuple[str, str], str] = {}  # (user_id, thread_id) -> sandbox_id
         self._acquire_serializer: AcquireSerializer[tuple[str, str]] = AcquireSerializer(thread_name_prefix="aio-sandbox-lock-wait")
+        # Lock waiters can block every serializer worker on one hot key. Work
+        # performed *after* a key is held must therefore use a separate pool:
+        # submitting it behind those waiters makes the holder wait for workers
+        # that are themselves waiting for the holder to release the key.
+        self._acquire_worker_executor = ThreadPoolExecutor(thread_name_prefix="aio-sandbox-owned-worker")
         self._last_activity: dict[str, float] = {}  # sandbox_id -> last activity timestamp
         # Warm pool: released sandboxes whose containers are still running.
         # Maps sandbox_id -> (SandboxInfo, release_timestamp).
@@ -2130,7 +2174,12 @@ class AioSandboxProvider(WarmPoolLifecycleMixin[SandboxInfo], SandboxProvider):
     async def _acquire_internal_async(self, thread_id: str | None, *, user_id: str) -> str:
         """Async counterpart to ``_acquire_internal``."""
         await asyncio.to_thread(self._ensure_skills_projection, user_id)
-        cached_id = await asyncio.to_thread(self._reuse_in_process_sandbox, thread_id, user_id=user_id)
+        cached_id = await _run_started_acquire_worker(
+            self._acquire_worker_executor,
+            self._reuse_in_process_sandbox,
+            thread_id,
+            user_id=user_id,
+        )
         if cached_id is not None:
             return cached_id
 
@@ -2142,7 +2191,13 @@ class AioSandboxProvider(WarmPoolLifecycleMixin[SandboxInfo], SandboxProvider):
                 self._assert_active_identity_available_locked(sandbox_id, key)
 
         # ── Layer 1.5: Warm pool (container still running, no cold-start) ──
-        reclaimed_id = await asyncio.to_thread(self._reclaim_warm_pool_sandbox, thread_id, sandbox_id, user_id=user_id)
+        reclaimed_id = await _run_started_acquire_worker(
+            self._acquire_worker_executor,
+            self._reclaim_warm_pool_sandbox,
+            thread_id,
+            sandbox_id,
+            user_id=user_id,
+        )
         if reclaimed_id is not None:
             return reclaimed_id
 
@@ -2598,6 +2653,7 @@ class AioSandboxProvider(WarmPoolLifecycleMixin[SandboxInfo], SandboxProvider):
     def reset(self) -> None:
         """Release process-local acquire workers when this instance is detached."""
         self._acquire_serializer.close()
+        self._acquire_worker_executor.shutdown(wait=False, cancel_futures=True)
 
     def shutdown(self) -> None:
         """Shutdown all sandboxes. Thread-safe and idempotent."""
@@ -2638,3 +2694,4 @@ class AioSandboxProvider(WarmPoolLifecycleMixin[SandboxInfo], SandboxProvider):
             logger.warning(f"Error closing sandbox ownership store during shutdown: {e}")
 
         self._acquire_serializer.close()
+        self._acquire_worker_executor.shutdown(wait=False, cancel_futures=True)

backend/packages/harness/deerflow/sandbox/AGENTS.md

@@ -18,7 +18,7 @@ cancellation before propagating it.
 **Implementations**:
 - `LocalSandboxProvider` - Local filesystem execution. `acquire(thread_id)` returns a per-user/thread `LocalSandbox` (id `local:{user_id}:{thread_id}`) whose `path_mappings` resolve `/mnt/user-data/{workspace,uploads,outputs}` and `/mnt/acp-workspace` to that thread's host directories, so the public `Sandbox` API honours the `/mnt/user-data` contract uniformly with AIO. `acquire()` / `acquire(None)` keeps the legacy generic singleton (id `local`) for callers without a thread context. Per-thread sandboxes are held in an LRU cache (default 256 entries) guarded by a `threading.Lock`. Shared runs use category mappings; a policy-scoped run replaces them with one `/mnt/skills` root mapping to the coherent thread view, so structured file tools resolve through one managed boundary. This is not a host filesystem security boundary: an enabled host `bash` subprocess can use canonical paths without `PathMapping`, so `supports_agent_skill_isolation` is dynamic and explicit Agent policies fail closed while host bash is enabled. Host-to-virtual output masking scans dynamic per-user/per-thread roots directly instead of compiling path-specific regexes, so evicted thread IDs do not remain in Python's global regex caches; a separate 256-entry root cache prevents repeated `realpath()` walks for every glob/grep match while bounding dynamic-path retention, and only the small process-stable skill/integration source set uses a bounded compiled cache. The shared `path_patterns.py` tail stops at `:`, so `$PATH`-style lists mask every entry; a mount symlink resolving outside all mounts keeps its mount path. On Windows, Git Bash/MSYS argument-conversion exclusions are limited to safe non-root virtual path prefixes; do not restore a blanket conversion disable, because host-native CLI launchers need normal MSYS path conversion for their own installation paths.
   PowerShell fallback pairs guarded console UTF-8 setters with explicit UTF-8 pipe decoding; cmd/MSYS retain locale decoding. Encoding regressions include real-pipe probes and Windows-only PowerShell roundtrips with and without a console.
-- `AioSandboxProvider` (`packages/harness/deerflow/community/`) - Docker isolation. Reuse enforces minimum persisted shell capacity even without image-env overrides. Upgrades require ownership-fenced teardown; provisioner create returns 409 for undersized Pods without deleting them. Acquire/reuse checks active-cache and warm-pool entries with the backend, dropping definitively dead containers from all in-process maps before rediscovery/creation. Health-check failures mean unknown, not dead; local discovery skips unverifiable containers and tries creation. Acquire/reclaim publishes ownership; `_renew_owned_leases` refreshes it on a dedicated thread. Text appends use native AIO append, never read-modify-write that could overwrite after a failed read. `reset()` closes the acquire serializer and its workers; `shutdown()` performs full remote teardown. Async create polls flock acquisition so cancelled waiters leave promptly. Once locked, cancellation drains admitted workers and create readiness/registration/cleanup before unlock/close. `uses_thread_data_mounts` defaults to backend detection (`LocalContainerBackend=True`, remote/provisioner=False); `sandbox.thread_data_mounts` overrides this for deployments sharing thread user-data directories. `true` skips upload acquire/sync; a false positive makes uploads unavailable. Explicit Agent policies use four thread projection category mounts and a distinct deterministic identity, excluding older shared-mount containers. The startup snapshot of `skills.container_path` drives mounts, identity, remote requests, and provisioner validation; root changes prevent container/Pod reuse. Gateway and provisioner independently require a canonical absolute root outside reserved platform mounts and derive four category allowlist entries from it. All four category overrides suppress the provisioner's default hostPath/skills-PVC mount; with `USERDATA_PVC_NAME`, categories use subpaths on that shared PVC. Readiness probes and `agent_sandbox` clients set `trust_env=False` for direct control-plane destinations (loopback/private IPs, single-label cluster hosts, Docker/Podman internal hostnames); external FQDNs and public IPs retain environment proxies.
+- `AioSandboxProvider` (`packages/harness/deerflow/community/`) - Docker isolation. Reuse enforces minimum persisted shell capacity even without image-env overrides. Upgrades require ownership-fenced teardown; provisioner create returns 409 for undersized Pods without deleting them. Acquire/reuse drops definitively dead active/warm entries before rediscovery/creation. Health-check failures mean unknown, so local discovery skips unverifiable containers. Acquire/reclaim publishes ownership; `_renew_owned_leases` refreshes it. Text appends use native AIO append, never read-modify-write. `reset()` closes acquire pools; `shutdown()` performs full teardown. Async acquire separates lock waiters from holder work; started cache/warm workers drain, queued work/skills projection stay cancellable. Create polls flock; owned work/readiness/registration/cleanup drain before unlock. `uses_thread_data_mounts` defaults to backend detection (`LocalContainerBackend=True`, remote/provisioner=False); `sandbox.thread_data_mounts` overrides this for deployments sharing thread user-data directories. `true` skips upload acquire/sync; a false positive makes uploads unavailable. Explicit Agent policies use four thread projection category mounts and a distinct deterministic identity, excluding older shared-mount containers. The startup snapshot of `skills.container_path` drives mounts, identity, remote requests, and provisioner validation; root changes prevent container/Pod reuse. Gateway and provisioner independently require a canonical absolute root outside reserved platform mounts and derive four category allowlist entries from it. All four category overrides suppress the provisioner's default hostPath/skills-PVC mount; with `USERDATA_PVC_NAME`, categories use subpaths on that shared PVC. Readiness probes and `agent_sandbox` clients set `trust_env=False` for direct control-plane destinations (loopback/private IPs, single-label cluster hosts, Docker/Podman internal hostnames); external FQDNs and public IPs retain environment proxies.
 - `E2BSandboxProvider` (`packages/harness/deerflow/community/e2b_sandbox/`) provides E2B remote isolation.
   New unrestricted sandboxes receive a one-shot upload from the enabled-only
   public, custom, legacy, and managed integration projections. For a thread

backend/tests/test_aio_sandbox_provider.py

@@ -296,6 +296,7 @@ def _make_provider(tmp_path):
         provider._acquire_epoch_counter = 0
         provider._acquire_inflight = {}
         provider._acquire_serializer = AcquireSerializer(thread_name_prefix="aio-sandbox-lock-wait")
+        provider._acquire_worker_executor = aio_mod.ThreadPoolExecutor(thread_name_prefix="aio-sandbox-owned-worker-test")
         provider._lock = MagicMock()
         provider._idle_checker_stop = MagicMock()
         provider._renewal_stop = MagicMock()
@@ -1374,6 +1375,244 @@ async def test_acquire_async_cancelled_waiter_does_not_block_successor(tmp_path,
     pytest.fail("provider thread lock was not released after successor acquire_async")
 
 
+@pytest.mark.anyio
+@pytest.mark.parametrize("blocked_step", ["reuse", "reclaim"])
+async def test_acquire_async_cancellation_keeps_serializer_until_started_worker_finishes(
+    blocked_step,
+    tmp_path,
+    monkeypatch,
+):
+    """A cancelled same-key acquire must not overlap an already-started worker."""
+    aio_mod = importlib.import_module("deerflow.community.aio_sandbox.aio_sandbox_provider")
+    provider = _make_provider(tmp_path)
+    provider._warm_pool = {}
+    provider._sandbox_infos = {}
+    provider._thread_sandboxes = {}
+    provider._sandboxes = {}
+    provider._last_activity = {}
+    provider._lock = aio_mod.threading.Lock()
+
+    first_started = threading.Event()
+    allow_first_finish = threading.Event()
+    successor_started = threading.Event()
+    state_lock = threading.Lock()
+    active = 0
+    max_active = 0
+    calls = 0
+
+    monkeypatch.setattr(provider, "_ensure_skills_projection", lambda _user_id: None)
+
+    def blocked_worker(result):
+        nonlocal active, max_active, calls
+        with state_lock:
+            calls += 1
+            call_number = calls
+            active += 1
+            max_active = max(max_active, active)
+        try:
+            if call_number == 1:
+                first_started.set()
+                assert allow_first_finish.wait(timeout=2)
+            else:
+                successor_started.set()
+            return result
+        finally:
+            with state_lock:
+                active -= 1
+
+    def reuse(*_args, **_kwargs):
+        if blocked_step == "reuse":
+            return blocked_worker("sandbox-reused")
+        return None
+
+    def reclaim(*_args, **_kwargs):
+        if blocked_step == "reclaim":
+            return blocked_worker("sandbox-reclaimed")
+        raise AssertionError("warm reclaim should not run after cached reuse succeeds")
+
+    monkeypatch.setattr(provider, "_reuse_in_process_sandbox", reuse)
+    monkeypatch.setattr(provider, "_reclaim_warm_pool_sandbox", reclaim)
+
+    owner = asyncio.create_task(provider.acquire_async("thread-owned-worker", user_id="default"))
+    successor = None
+    try:
+        assert await asyncio.to_thread(first_started.wait, 2)
+        owner.cancel()
+        successor = asyncio.create_task(provider.acquire_async("thread-owned-worker", user_id="default"))
+
+        await asyncio.sleep(0.05)
+        assert not successor_started.is_set(), "serializer released while cancelled worker was still running"
+
+        allow_first_finish.set()
+        with pytest.raises(asyncio.CancelledError):
+            await owner
+
+        assert await asyncio.to_thread(successor_started.wait, 2)
+        expected = "sandbox-reused" if blocked_step == "reuse" else "sandbox-reclaimed"
+        assert await asyncio.wait_for(successor, timeout=1) == expected
+        assert max_active == 1
+    finally:
+        allow_first_finish.set()
+        if not owner.done():
+            owner.cancel()
+            await asyncio.gather(owner, return_exceptions=True)
+        if successor is not None and not successor.done():
+            successor.cancel()
+            await asyncio.gather(successor, return_exceptions=True)
+        provider.reset()
+
+
+@pytest.mark.anyio
+async def test_acquire_async_cancellation_cancels_queued_reclaim_before_it_runs(tmp_path, monkeypatch):
+    """Cancellation must not force a not-yet-started reclaim to run later."""
+    aio_mod = importlib.import_module("deerflow.community.aio_sandbox.aio_sandbox_provider")
+    provider = _make_provider(tmp_path)
+    provider._acquire_worker_executor.shutdown(wait=False, cancel_futures=True)
+    provider._acquire_worker_executor = aio_mod.ThreadPoolExecutor(
+        max_workers=1,
+        thread_name_prefix="aio-owned-worker-test",
+    )
+    provider._warm_pool = {}
+    provider._sandbox_infos = {}
+    provider._thread_sandboxes = {}
+    provider._sandboxes = {}
+    provider._last_activity = {}
+    provider._lock = aio_mod.threading.Lock()
+
+    blocker_started = threading.Event()
+    allow_blocker = threading.Event()
+    reclaim_submitted = threading.Event()
+    reclaim_calls = 0
+
+    executor = provider._acquire_worker_executor
+    original_submit = executor.submit
+    submit_count = 0
+
+    def occupy_lifecycle_executor():
+        blocker_started.set()
+        assert allow_blocker.wait(timeout=2)
+
+    def reuse(*_args, **_kwargs):
+        # We are running on the owned-lifecycle executor. Queue the blocker
+        # behind this worker before returning None; with max_workers=1 it starts
+        # immediately after reuse returns and before the event loop can submit
+        # warm reclaim.
+        original_submit(occupy_lifecycle_executor)
+        return None
+
+    def reclaim(*_args, **_kwargs):
+        nonlocal reclaim_calls
+        reclaim_calls += 1
+        return "sandbox-reclaimed"
+
+    def tracking_submit(*args, **kwargs):
+        nonlocal submit_count
+        submit_count += 1
+        future = original_submit(*args, **kwargs)
+        # Submission 1 runs cached reuse; submission 2 is warm reclaim,
+        # now queued behind occupy_lifecycle_executor.
+        if submit_count == 2:
+            reclaim_submitted.set()
+        return future
+
+    monkeypatch.setattr(provider, "_ensure_skills_projection", lambda _user_id: None)
+    monkeypatch.setattr(provider, "_reuse_in_process_sandbox", reuse)
+    monkeypatch.setattr(provider, "_reclaim_warm_pool_sandbox", reclaim)
+    monkeypatch.setattr(executor, "submit", tracking_submit)
+
+    owner = asyncio.create_task(provider.acquire_async("thread-queued-reclaim", user_id="default"))
+    try:
+        assert await asyncio.to_thread(blocker_started.wait, 2)
+        assert await asyncio.to_thread(reclaim_submitted.wait, 2)
+
+        owner.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await owner
+
+        allow_blocker.set()
+        await asyncio.sleep(0.05)
+        assert reclaim_calls == 0
+    finally:
+        allow_blocker.set()
+        if not owner.done():
+            owner.cancel()
+            await asyncio.gather(owner, return_exceptions=True)
+        provider._acquire_serializer.close()
+        provider._acquire_worker_executor.shutdown(wait=False, cancel_futures=True)
+
+
+@pytest.mark.anyio
+async def test_acquire_async_lock_waiters_do_not_starve_holder_worker(tmp_path, monkeypatch):
+    """Lock waiters must not occupy the executor needed by the lock holder."""
+    provider = _make_provider(tmp_path)
+    provider._acquire_serializer.close()
+    provider._acquire_serializer = AcquireSerializer(
+        max_workers=2,
+        thread_name_prefix="aio-holder-deadlock-test",
+    )
+
+    projection_started = threading.Event()
+    allow_projection = threading.Event()
+    waiter_workers_started = threading.Event()
+    count_lock = threading.Lock()
+    started_workers = 0
+    projection_calls = 0
+
+    executor = provider._acquire_serializer.executor
+    original_submit = executor.submit
+
+    def tracking_submit(func, /, *args, **kwargs):
+        def tracked():
+            nonlocal started_workers
+            with count_lock:
+                started_workers += 1
+                if started_workers >= 3:
+                    waiter_workers_started.set()
+            return func(*args, **kwargs)
+
+        return original_submit(tracked)
+
+    def ensure_projection(_user_id):
+        nonlocal projection_calls
+        projection_calls += 1
+        if projection_calls == 1:
+            projection_started.set()
+            assert allow_projection.wait(timeout=2)
+
+    monkeypatch.setattr(executor, "submit", tracking_submit)
+    monkeypatch.setattr(provider, "_ensure_skills_projection", ensure_projection)
+    monkeypatch.setattr(
+        provider,
+        "_reuse_in_process_sandbox",
+        lambda *_args, **_kwargs: "sandbox-cached",
+    )
+
+    owner = asyncio.create_task(provider.acquire_async("thread-holder-deadlock", user_id="default"))
+    successors: list[asyncio.Task[str]] = []
+    try:
+        assert await asyncio.to_thread(projection_started.wait, 2)
+        successors = [asyncio.create_task(provider.acquire_async("thread-holder-deadlock", user_id="default")) for _ in range(2)]
+        # Submission 1 acquired the holder's key. Submissions 2 and 3 are now
+        # both running in the bounded serializer pool, blocked on that same key.
+        assert await asyncio.to_thread(waiter_workers_started.wait, 2)
+
+        allow_projection.set()
+        assert await asyncio.wait_for(owner, timeout=1) == "sandbox-cached"
+        assert await asyncio.wait_for(
+            asyncio.gather(*successors),
+            timeout=1,
+        ) == ["sandbox-cached", "sandbox-cached"]
+    finally:
+        allow_projection.set()
+        if not owner.done():
+            owner.cancel()
+        for task in successors:
+            if not task.done():
+                task.cancel()
+        await asyncio.gather(owner, *successors, return_exceptions=True)
+        provider.reset()
+
+
 @pytest.mark.anyio
 async def test_acquire_internal_async_offloads_cached_reuse_health_check(tmp_path, monkeypatch):
     """Async cached reuse must keep backend health checks off the event loop."""
@@ -1395,7 +1634,6 @@ async def test_acquire_internal_async_offloads_cached_reuse_health_check(tmp_pat
     assert sandbox_id == "sandbox-cached-async"
     assert to_thread_calls == [
         (provider._ensure_skills_projection, ("default",)),
-        (provider._reuse_in_process_sandbox, ("thread-cached-async",)),
     ]
 
 

backend/tests/test_sandbox_orphan_reconciliation.py

@@ -434,6 +434,7 @@ def _make_provider_for_reconciliation(tmp_path=None, *, worker_id: str = "worker
     provider._sandbox_infos = {}
     provider._thread_sandboxes = {}
     provider._acquire_serializer = AcquireSerializer(thread_name_prefix="aio-sandbox-lock-wait")
+    provider._acquire_worker_executor = aio_mod.ThreadPoolExecutor(thread_name_prefix="aio-sandbox-owned-worker-test")
     provider._last_activity = {}
     provider._warm_pool = {}
     provider._active_sandbox_identity = {}