"""Execution-scoped leases for process-local sandbox use. Provider ownership stores answer which Gateway instance may reap a remote sandbox. This module answers a different question: which concurrently running agent executions inside one Gateway are still using the provider's active client. The last execution lease is the only one allowed to call ``SandboxProvider.release``. """ from __future__ import annotations import asyncio import logging import threading import uuid from collections.abc import Callable from dataclasses import dataclass from typing import TYPE_CHECKING, Any from deerflow.sandbox.acquire_serialization import AcquireSerializer if TYPE_CHECKING: from deerflow.sandbox.sandbox_provider import SandboxProvider logger = logging.getLogger(__name__) SANDBOX_LEASE_OWNER_CONTEXT_KEY = "sandbox_lease_owner_id" SANDBOX_COMMAND_SCOPE_CONTEXT_KEY = "sandbox_command_scope_id" SANDBOX_SERVER_OWNED_CONTEXT_KEYS = frozenset( { SANDBOX_LEASE_OWNER_CONTEXT_KEY, SANDBOX_COMMAND_SCOPE_CONTEXT_KEY, } ) async def _drain_task_after_cancellation[T](task: asyncio.Task[T]) -> T: """Wait for ``task`` even if the current task is cancelled again. Lifecycle reconciliation must keep its serializer until provider work has finished. Repeated cancellation is therefore remembered by the caller but cannot propagate into the reconciliation task or interrupt this drain. """ while True: try: return await asyncio.shield(task) except asyncio.CancelledError: if task.done(): return task.result() async def run_sync_lifecycle_operation[T](func: Callable[..., T], /, *args: Any, **kwargs: Any) -> T: """Run blocking client work without letting cancellation outlive it. ``asyncio.to_thread`` cannot stop its worker when the awaiting task is cancelled. Sandbox cleanup must therefore wait for the worker before an outer execution/request fence is allowed to release its client holder. Repeated cancellation is remembered and propagated only after the worker has terminated; a late worker failure is logged without replacing the caller's cancellation. """ operation_task = asyncio.create_task(asyncio.to_thread(func, *args, **kwargs)) try: return await asyncio.shield(operation_task) except asyncio.CancelledError as cancellation: try: await _drain_task_after_cancellation(operation_task) except Exception: logger.warning( "Cancelled sandbox client operation failed while draining", exc_info=True, ) raise cancellation @dataclass(slots=True) class SandboxClientLease: """One bounded caller's process-local hold on a sandbox client.""" sandbox: object | None sandbox_id: str owner_id: str | None provider: SandboxProvider async def release(self) -> None: """Drop this holder after its final sandbox operation has drained.""" if self.owner_id is None: return owner_id = self.owner_id self.owner_id = None await get_sandbox_lease_manager(self.provider).release_async(owner_id) async def run_sync[T](self, func: Callable[..., T], /, *args: Any, **kwargs: Any) -> T: """Run one blocking client operation inside this lease boundary.""" return await run_sync_lifecycle_operation(func, *args, **kwargs) async def acquire_sandbox_client_lease( provider: SandboxProvider, thread_id: str, *, user_id: str, owner_prefix: str, release_on_last: bool = True, ) -> SandboxClientLease: """Acquire a unique holder and resolve its process-local client. This is the non-HTTP counterpart of Gateway request leases. Callers must keep the returned object through their last sandbox operation and release it in ``finally``. ``release_on_last=False`` is appropriate for upload synchronization: the upload fences a concurrent run but does not itself request that the warm sandbox be parked. """ owner_id = f"{owner_prefix}:{uuid.uuid4()}" manager = get_sandbox_lease_manager(provider) sandbox_id = await manager.acquire_async( owner_id, thread_id, user_id=user_id, release_on_last=release_on_last, ) try: sandbox = provider.get(sandbox_id) except BaseException: await manager.release_async(owner_id) raise return SandboxClientLease( sandbox=sandbox, sandbox_id=sandbox_id, owner_id=owner_id, provider=provider, ) @dataclass(frozen=True) class _LeaseBinding: sandbox_id: str thread_key: tuple[str, str] release_on_last: bool = True class SandboxLeaseManager: """Coordinate active agent users of one sandbox provider. Lifecycle transitions are serialized per user/thread key. Metadata is protected separately so unrelated threads do not block each other's slow provider operations. """ def __init__(self, provider: SandboxProvider): self._provider = provider self._metadata_lock = threading.RLock() self._serializer = AcquireSerializer[tuple[str, str]]( thread_name_prefix="sandbox-execution-lease", ) self._bindings_by_owner: dict[str, _LeaseBinding] = {} self._owners_by_sandbox: dict[str, set[str]] = {} self._release_pending_by_sandbox: set[str] = set() @staticmethod def _thread_key(thread_id: str, user_id: str) -> tuple[str, str]: return user_id, thread_id def _remove_binding_locked( self, owner_id: str, *, request_release: bool = True, ) -> tuple[_LeaseBinding | None, bool]: binding = self._bindings_by_owner.pop(owner_id, None) if binding is None: return None, False if request_release and binding.release_on_last: # Closing responsibility belongs to the sandbox lifecycle, not to # whichever holder happens to finish last. A normal owner may # leave while a fork/upload borrower is still using the client; # remember its close request until every holder is gone. self._release_pending_by_sandbox.add(binding.sandbox_id) owners = self._owners_by_sandbox.get(binding.sandbox_id) if owners is None: release_provider = binding.sandbox_id in self._release_pending_by_sandbox self._release_pending_by_sandbox.discard(binding.sandbox_id) return binding, release_provider owners.discard(owner_id) if owners: return binding, False self._owners_by_sandbox.pop(binding.sandbox_id, None) release_provider = binding.sandbox_id in self._release_pending_by_sandbox self._release_pending_by_sandbox.discard(binding.sandbox_id) return binding, release_provider def _bind_locked( self, owner_id: str, sandbox_id: str, key: tuple[str, str], *, release_on_last: bool, ) -> tuple[_LeaseBinding | None, bool]: existing = self._bindings_by_owner.get(owner_id) if existing is not None and existing.thread_key != key: raise RuntimeError(f"Sandbox lease owner {owner_id!r} cannot move between thread identities") if existing is not None and existing.sandbox_id == sandbox_id: # Ownership is monotonic for one execution. A normal owner may # later encounter a fork-restored view, but must not lose its # responsibility to park the provider. A borrowed owner can be # upgraded when it later performs a normal acquisition. if release_on_last and not existing.release_on_last: self._bindings_by_owner[owner_id] = _LeaseBinding( sandbox_id=sandbox_id, thread_key=key, release_on_last=True, ) return None, False release_previous = False previous: _LeaseBinding | None = None if existing is not None: previous, release_previous = self._remove_binding_locked(owner_id) self._bindings_by_owner[owner_id] = _LeaseBinding( sandbox_id=sandbox_id, thread_key=key, release_on_last=release_on_last, ) self._owners_by_sandbox.setdefault(sandbox_id, set()).add(owner_id) return previous, release_previous def _release_unbound_acquire(self, sandbox_id: str) -> None: """Undo a cancelled acquire when no admitted execution uses its result. The caller must still hold the serializer for the originating thread key. That makes the owner check and provider release one transition: another execution for the same thread cannot bind the sandbox between them. """ with self._metadata_lock: has_owners = bool(self._owners_by_sandbox.get(sandbox_id)) if not has_owners: self._provider.release(sandbox_id) def _active_owner_binding( self, owner_id: str, key: tuple[str, str], *, release_on_last: bool, ) -> str | None: """Return one live owner binding while the caller holds ``key``. Middleware may retain a checkpointed id before the provider discovers that its local client is gone. Treat that owner binding as stale so a later acquire rebuilds it instead of preserving an unusable id. """ with self._metadata_lock: existing = self._bindings_by_owner.get(owner_id) if existing is None: return None if existing.thread_key != key: raise RuntimeError(f"Sandbox lease owner {owner_id!r} cannot move between thread identities") if self._provider.get(existing.sandbox_id) is not None: if release_on_last and not existing.release_on_last: with self._metadata_lock: if self._bindings_by_owner.get(owner_id) == existing: self._bindings_by_owner[owner_id] = _LeaseBinding( sandbox_id=existing.sandbox_id, thread_key=existing.thread_key, release_on_last=True, ) return existing.sandbox_id with self._metadata_lock: if self._bindings_by_owner.get(owner_id) == existing: self._remove_binding_locked(owner_id, request_release=False) return None def _acquire_and_bind( self, owner_id: str, thread_id: str, user_id: str, key: tuple[str, str], *, release_on_last: bool, ) -> str: sandbox_id = self._provider.acquire(thread_id, user_id=user_id) with self._metadata_lock: previous, release_previous = self._bind_locked( owner_id, sandbox_id, key, release_on_last=release_on_last, ) if release_previous and previous is not None: self._provider.release(previous.sandbox_id) return sandbox_id async def _acquire_and_bind_async( self, owner_id: str, thread_id: str, user_id: str, key: tuple[str, str], *, release_on_last: bool, ) -> str: acquire_task = asyncio.create_task(self._provider.acquire_async(thread_id, user_id=user_id)) try: sandbox_id = await asyncio.shield(acquire_task) except asyncio.CancelledError as cancellation: # Provider implementations commonly offload container startup to a # worker thread, which cannot be stopped by cancelling the awaiter. # Reconcile the result before relinquishing the serializer. try: sandbox_id = await _drain_task_after_cancellation(acquire_task) except Exception: logger.warning( "Cancelled sandbox acquire failed during reconciliation", exc_info=True, ) raise cancellation rollback_task = asyncio.create_task( asyncio.to_thread( self._release_unbound_acquire, sandbox_id, ) ) try: await _drain_task_after_cancellation(rollback_task) except Exception: logger.warning( "Cancelled sandbox acquire rollback failed during reconciliation", exc_info=True, ) raise cancellation with self._metadata_lock: previous, release_previous = self._bind_locked( owner_id, sandbox_id, key, release_on_last=release_on_last, ) if release_previous and previous is not None: await asyncio.to_thread( self._provider.release, previous.sandbox_id, ) return sandbox_id def retain( self, owner_id: str, sandbox_id: str, *, thread_id: str, user_id: str, release_on_last: bool = True, ) -> None: """Attach an execution to an inherited or checkpointed sandbox id. ``release_on_last=False`` is for a borrower: it fences the client and owns command-scope cleanup without requesting a park itself. A normal owner's earlier park request can still be deferred until it drains. """ key = self._thread_key(thread_id, user_id) with self._serializer.hold(key): with self._metadata_lock: previous, release_previous = self._bind_locked( owner_id, sandbox_id, key, release_on_last=release_on_last, ) if release_previous and previous is not None: self._provider.release(previous.sandbox_id) async def retain_async( self, owner_id: str, sandbox_id: str, *, thread_id: str, user_id: str, release_on_last: bool = True, ) -> None: """Async attach without blocking the event loop on a lifecycle lock.""" key = self._thread_key(thread_id, user_id) async with self._serializer.hold_async(key): with self._metadata_lock: previous, release_previous = self._bind_locked( owner_id, sandbox_id, key, release_on_last=release_on_last, ) if release_previous and previous is not None: await asyncio.to_thread( self._provider.release, previous.sandbox_id, ) def acquire( self, owner_id: str, thread_id: str, *, user_id: str, release_on_last: bool = True, ) -> str: """Acquire and bind a sandbox, idempotently for one execution owner.""" key = self._thread_key(thread_id, user_id) with self._serializer.hold(key): existing_sandbox_id = self._active_owner_binding( owner_id, key, release_on_last=release_on_last, ) if existing_sandbox_id is not None: return existing_sandbox_id return self._acquire_and_bind( owner_id, thread_id, user_id, key, release_on_last=release_on_last, ) def reuse_or_acquire( self, owner_id: str, sandbox_id: str, *, thread_id: str, user_id: str, release_on_last: bool = True, acquire_release_on_last: bool = True, ) -> str: """Atomically retain a live persisted sandbox or acquire a replacement. A fork-restored live client is borrowed with ``release_on_last=False``; if that persisted client is gone, its freshly acquired replacement is owned normally unless ``acquire_release_on_last`` is also disabled. """ key = self._thread_key(thread_id, user_id) with self._serializer.hold(key): existing_sandbox_id = self._active_owner_binding( owner_id, key, release_on_last=release_on_last, ) if existing_sandbox_id is not None: return existing_sandbox_id if self._provider.get(sandbox_id) is not None: with self._metadata_lock: previous, release_previous = self._bind_locked( owner_id, sandbox_id, key, release_on_last=release_on_last, ) if release_previous and previous is not None: self._provider.release(previous.sandbox_id) return sandbox_id return self._acquire_and_bind( owner_id, thread_id, user_id, key, release_on_last=acquire_release_on_last, ) async def acquire_async( self, owner_id: str, thread_id: str, *, user_id: str, release_on_last: bool = True, ) -> str: """Async acquire while preserving the provider's own async hook.""" key = self._thread_key(thread_id, user_id) async with self._serializer.hold_async(key): existing_sandbox_id = self._active_owner_binding( owner_id, key, release_on_last=release_on_last, ) if existing_sandbox_id is not None: return existing_sandbox_id return await self._acquire_and_bind_async( owner_id, thread_id, user_id, key, release_on_last=release_on_last, ) async def reuse_or_acquire_async( self, owner_id: str, sandbox_id: str, *, thread_id: str, user_id: str, release_on_last: bool = True, acquire_release_on_last: bool = True, ) -> str: """Async atomic retain-or-replace transition for a persisted sandbox.""" key = self._thread_key(thread_id, user_id) async with self._serializer.hold_async(key): existing_sandbox_id = self._active_owner_binding( owner_id, key, release_on_last=release_on_last, ) if existing_sandbox_id is not None: return existing_sandbox_id if self._provider.get(sandbox_id) is not None: with self._metadata_lock: previous, release_previous = self._bind_locked( owner_id, sandbox_id, key, release_on_last=release_on_last, ) if release_previous and previous is not None: await asyncio.to_thread( self._provider.release, previous.sandbox_id, ) return sandbox_id return await self._acquire_and_bind_async( owner_id, thread_id, user_id, key, release_on_last=acquire_release_on_last, ) def release(self, owner_id: str) -> None: """Release one execution and park the sandbox only after the last user.""" with self._metadata_lock: binding = self._bindings_by_owner.get(owner_id) if binding is None: return with self._serializer.hold(binding.thread_key): with self._metadata_lock: current = self._bindings_by_owner.get(owner_id) if current is None: return binding, release_provider = self._remove_binding_locked(owner_id) assert binding is not None try: sandbox = self._provider.get(binding.sandbox_id) if sandbox is not None: sandbox.release_command_scope(owner_id) finally: if release_provider: self._provider.release(binding.sandbox_id) async def release_async(self, owner_id: str) -> None: """Release a lease without blocking the caller's event loop.""" release_task = asyncio.create_task(asyncio.to_thread(self.release, owner_id)) try: await asyncio.shield(release_task) except asyncio.CancelledError as cancellation: # Complete lifecycle cleanup before allowing cancellation to leave # the agent's finally block. try: await _drain_task_after_cancellation(release_task) except Exception: logger.warning( "Cancelled sandbox release failed during reconciliation", exc_info=True, ) raise cancellation raise cancellation def binding_for(self, owner_id: str) -> str | None: """Return the sandbox bound to an owner; intended for diagnostics/tests.""" with self._metadata_lock: binding = self._bindings_by_owner.get(owner_id) return binding.sandbox_id if binding is not None else None def close(self) -> None: """Stop accepting new transitions and release serializer workers.""" self._serializer.close() _manager_lock = threading.Lock() _managers: dict[int, tuple[SandboxProvider, SandboxLeaseManager]] = {} def get_sandbox_lease_manager(provider: SandboxProvider) -> SandboxLeaseManager: """Return the process-local lease manager for this provider object. Provider implementations are not required to be hashable, and distinct instances that compare equal must not share lifecycle state. Keep a strong identity entry until the provider is explicitly detached; the manager already owns the provider strongly, so a weak-key registry would not make the lifecycle shorter. """ provider_id = id(provider) with _manager_lock: entry = _managers.get(provider_id) if entry is not None and entry[0] is provider: return entry[1] manager = SandboxLeaseManager(provider) _managers[provider_id] = (provider, manager) return manager def discard_sandbox_lease_manager(provider: SandboxProvider) -> None: """Forget lease metadata when a provider singleton is detached.""" provider_id = id(provider) with _manager_lock: entry = _managers.get(provider_id) if entry is None or entry[0] is not provider: manager = None else: _, manager = _managers.pop(provider_id) if manager is not None: manager.close() def ensure_sandbox_lease_owner(context: Any) -> str | None: """Create one ephemeral owner id in a mutable runtime context.""" if not isinstance(context, dict): return None existing = context.get(SANDBOX_LEASE_OWNER_CONTEXT_KEY) if isinstance(existing, str) and existing: return existing owner_id = f"agent:{uuid.uuid4()}" context[SANDBOX_LEASE_OWNER_CONTEXT_KEY] = owner_id return owner_id def sandbox_lease_owner(context: Any) -> str | None: """Read an execution owner without creating one for direct tool callers.""" if not isinstance(context, dict): return None owner_id = context.get(SANDBOX_LEASE_OWNER_CONTEXT_KEY) return owner_id if isinstance(owner_id, str) and owner_id else None def sandbox_command_scope(context: Any) -> str | None: """Read the optional shell-session scope carried by subagent execution.""" if not isinstance(context, dict): return None scope_id = context.get(SANDBOX_COMMAND_SCOPE_CONTEXT_KEY) return scope_id if isinstance(scope_id, str) and scope_id else None def _sandbox_execution_lease(context: Any) -> tuple[str, str] | None: """Return a bound execution lease without initializing a provider.""" owner_id = sandbox_lease_owner(context) if owner_id is None or not isinstance(context, dict): return None sandbox_id = context.get("sandbox_id") if not isinstance(sandbox_id, str) or not sandbox_id: return None return owner_id, sandbox_id def release_sandbox_execution_lease(context: Any) -> None: """Release a lead/embedded execution lease at its outer lifecycle fence.""" lease = _sandbox_execution_lease(context) if lease is None: return # Import lazily so runs that never touch a sandbox do not initialize the # provider during terminal cleanup. from deerflow.sandbox.sandbox_provider import get_sandbox_provider owner_id, _sandbox_id = lease provider = get_sandbox_provider() get_sandbox_lease_manager(provider).release(owner_id) async def release_sandbox_execution_lease_async(context: Any) -> None: """Async counterpart to :func:`release_sandbox_execution_lease`.""" lease = _sandbox_execution_lease(context) if lease is None: return from deerflow.sandbox.sandbox_provider import get_sandbox_provider owner_id, _sandbox_id = lease provider = get_sandbox_provider() await get_sandbox_lease_manager(provider).release_async(owner_id)