mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-11-08 03:30:34 +00:00
a17e5694cc
* fix(sandbox): stop owned command processes when a run is cancelled * fix(sandbox): handle incomplete AIO output responses --------- Co-authored-by: Totoro-qaq <279883115+Totoro-qaq@users.noreply.github.com> Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
96 lines
2.8 KiB
Python
96 lines
2.8 KiB
Python
"""Process ownership for cancellable bash calls; retained only for one execution lease."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import uuid
|
|
from contextlib import contextmanager
|
|
from contextvars import ContextVar
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
COMMAND_OWNER_ENV = "DEERFLOW_COMMAND_OWNER"
|
|
|
|
|
|
class CommandScopeCancelled(RuntimeError):
|
|
"""An explicitly cancelled execution must not start or retry commands."""
|
|
|
|
|
|
@dataclass
|
|
class CommandScope:
|
|
token: str = field(default_factory=lambda: uuid.uuid4().hex)
|
|
cancelled: threading.Event = field(default_factory=threading.Event)
|
|
cleaned: threading.Event = field(default_factory=threading.Event)
|
|
lock: threading.Lock = field(default_factory=threading.Lock)
|
|
processes: list[Any] = field(default_factory=list)
|
|
|
|
def check(self) -> None:
|
|
if self.cancelled.is_set():
|
|
raise CommandScopeCancelled("Sandbox execution was cancelled")
|
|
|
|
|
|
class CommandScopes:
|
|
def __init__(self) -> None:
|
|
self._lock = threading.Lock()
|
|
self._scopes: dict[str, CommandScope] = {}
|
|
|
|
def get(self, owner: str) -> CommandScope:
|
|
with self._lock:
|
|
scope = self._scopes.get(owner)
|
|
if scope is None:
|
|
scope = self._scopes[owner] = CommandScope()
|
|
return scope
|
|
|
|
def find(self, owner: str) -> CommandScope | None:
|
|
with self._lock:
|
|
return self._scopes.get(owner)
|
|
|
|
def request_cancel(self, owner: str) -> None:
|
|
scope = self.get(owner)
|
|
scope.cancelled.set()
|
|
|
|
def release(self, owner: str) -> None:
|
|
with self._lock:
|
|
self._scopes.pop(owner, None)
|
|
|
|
|
|
_active: ContextVar[tuple[object, CommandScope] | None] = ContextVar("sandbox_command_owner", default=None)
|
|
|
|
|
|
@contextmanager
|
|
def command_owner(sandbox: object, owner: str | None):
|
|
"""Bind server-owned execution identity across the synchronous tool bridge."""
|
|
if not owner or not getattr(sandbox, "supports_command_cancellation", False):
|
|
yield
|
|
return
|
|
scope = sandbox._command_scopes.get(owner)
|
|
scope.check()
|
|
token = _active.set((sandbox, scope))
|
|
try:
|
|
yield
|
|
finally:
|
|
_active.reset(token)
|
|
|
|
|
|
def current_command_scope(sandbox: object | None = None) -> CommandScope | None:
|
|
active = _active.get()
|
|
if active is None or (sandbox is not None and active[0] is not sandbox):
|
|
return None
|
|
return active[1]
|
|
|
|
|
|
@contextmanager
|
|
def command_lock(lock, scope: CommandScope | None):
|
|
"""A cancelled queued command must not wait for another execution's I/O."""
|
|
if scope is None:
|
|
with lock:
|
|
yield
|
|
return
|
|
while not lock.acquire(timeout=0.1):
|
|
scope.check()
|
|
try:
|
|
scope.check()
|
|
yield
|
|
finally:
|
|
lock.release()
|