diff --git a/backend/packages/harness/deerflow/sandbox/middleware.py b/backend/packages/harness/deerflow/sandbox/middleware.py index 9937bcdb6..36da18bfd 100644 --- a/backend/packages/harness/deerflow/sandbox/middleware.py +++ b/backend/packages/harness/deerflow/sandbox/middleware.py @@ -9,11 +9,12 @@ from langchain.agents.middleware import AgentMiddleware from langchain_core.messages import ToolMessage from langgraph.prebuilt.tool_node import ToolCallRequest from langgraph.runtime import Runtime -from langgraph.types import Command, Overwrite +from langgraph.types import Command from deerflow.agents.thread_state import SandboxStateField, ThreadDataState from deerflow.runtime.user_context import resolve_runtime_user_id from deerflow.sandbox import get_sandbox_provider +from deerflow.sandbox.overwrite import unwrap_sandbox logger = logging.getLogger(__name__) @@ -25,24 +26,6 @@ class SandboxMiddlewareState(AgentState): thread_data: NotRequired[ThreadDataState | None] -def _unwrap_sandbox(sandbox: object) -> tuple[object, bool]: - """Unwrap an ``Overwrite``-wrapped sandbox channel value, if present. - - Fork-restored checkpoints can deliver the sandbox channel still wrapped in - ``langgraph.types.Overwrite`` (the rollback restore applies replace-style - writes through a state-mutation graph in delta checkpoint mode). Reading - ``sandbox["sandbox_id"]`` on the wrapper itself crashes with ``TypeError: - 'Overwrite' object is not subscriptable``, so unwrap before use. - - Returns ``(value, fork_restored)``. The wrapped form replays the parent - thread's sandbox state, so callers must not treat the sandbox as owned by - this run (e.g. release it). - """ - if isinstance(sandbox, Overwrite): - return sandbox.value, True - return sandbox, False - - class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]): """Create a sandbox environment and assign it to an agent. @@ -117,7 +100,7 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]): @override def after_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None: - sandbox, fork_restored = _unwrap_sandbox(state.get("sandbox")) + sandbox, fork_restored = unwrap_sandbox(state.get("sandbox")) if sandbox is not None: sandbox_id = sandbox["sandbox_id"] if fork_restored: @@ -140,7 +123,7 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]): @override async def aafter_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None: - sandbox, fork_restored = _unwrap_sandbox(state.get("sandbox")) + sandbox, fork_restored = unwrap_sandbox(state.get("sandbox")) if sandbox is not None: sandbox_id = sandbox["sandbox_id"] if fork_restored: diff --git a/backend/packages/harness/deerflow/sandbox/overwrite.py b/backend/packages/harness/deerflow/sandbox/overwrite.py new file mode 100644 index 000000000..c9bea7bda --- /dev/null +++ b/backend/packages/harness/deerflow/sandbox/overwrite.py @@ -0,0 +1,21 @@ +"""Helpers for ``langgraph.types.Overwrite``-wrapped channel values.""" + +from langgraph.types import Overwrite + + +def unwrap_sandbox(sandbox: object) -> tuple[object, bool]: + """Unwrap an ``Overwrite``-wrapped sandbox channel value, if present. + + Fork-restored checkpoints can deliver the sandbox channel still wrapped in + ``langgraph.types.Overwrite`` (the rollback restore applies replace-style + writes through a state-mutation graph in delta checkpoint mode). Reading + ``sandbox["sandbox_id"]`` or ``sandbox.get("sandbox_id")`` on the wrapper + itself crashes, so unwrap before use. + + Returns ``(value, fork_restored)``. The wrapped form replays the parent + thread's sandbox state, so callers must not treat the sandbox as owned by + this run (e.g. release it). + """ + if isinstance(sandbox, Overwrite): + return sandbox.value, True + return sandbox, False diff --git a/backend/packages/harness/deerflow/sandbox/tools.py b/backend/packages/harness/deerflow/sandbox/tools.py index 9bcdc7c2b..6a7f6e32b 100644 --- a/backend/packages/harness/deerflow/sandbox/tools.py +++ b/backend/packages/harness/deerflow/sandbox/tools.py @@ -23,6 +23,7 @@ from deerflow.sandbox.exceptions import ( SandboxRuntimeError, ) from deerflow.sandbox.file_operation_lock import get_file_operation_lock +from deerflow.sandbox.overwrite import unwrap_sandbox from deerflow.sandbox.path_patterns import build_output_mask_pattern from deerflow.sandbox.sandbox import Sandbox from deerflow.sandbox.sandbox_provider import get_sandbox_provider @@ -1329,7 +1330,9 @@ def is_local_sandbox(runtime: Runtime | None) -> bool: return False if runtime.state is None: return False - sandbox_state = runtime.state.get("sandbox") + # Read-only classification: the id is only matched, so a fork-restored + # wrapper is safe to discard here (nothing gets released on this path). + sandbox_state, _ = unwrap_sandbox(runtime.state.get("sandbox")) if sandbox_state is None: return False sandbox_id = sandbox_state.get("sandbox_id") @@ -1352,7 +1355,9 @@ def sandbox_from_runtime(runtime: Runtime | None = None) -> Sandbox: raise SandboxRuntimeError("Tool runtime not available") if runtime.state is None: raise SandboxRuntimeError("Tool runtime state not available") - sandbox_state = runtime.state.get("sandbox") + # Read-only lookup: this only resolves the provider entry, and ownership + # (release) stays with after_agent's short-circuit on the wrapped state. + sandbox_state, _ = unwrap_sandbox(runtime.state.get("sandbox")) if sandbox_state is None: raise SandboxRuntimeError("Sandbox state not initialized in runtime") sandbox_id = sandbox_state.get("sandbox_id") @@ -1392,7 +1397,10 @@ def ensure_sandbox_initialized(runtime: Runtime | None = None) -> Sandbox: raise SandboxRuntimeError("Tool runtime state not available") # Check if sandbox already exists in state - sandbox_state = runtime.state.get("sandbox") + # Discarding fork_restored is safe: after_agent short-circuits on the + # still-wrapped state before the context-based release branch, so this + # reuse path never releases the parent sandbox. + sandbox_state, _ = unwrap_sandbox(runtime.state.get("sandbox")) if sandbox_state is not None: sandbox_id = sandbox_state.get("sandbox_id") if sandbox_id is not None: @@ -1439,7 +1447,9 @@ async def ensure_sandbox_initialized_async(runtime: Runtime | None = None) -> Sa if runtime.state is None: raise SandboxRuntimeError("Tool runtime state not available") - sandbox_state = runtime.state.get("sandbox") + # Same discard as the sync path above: the reuse path never releases, + # because after_agent short-circuits on the still-wrapped state first. + sandbox_state, _ = unwrap_sandbox(runtime.state.get("sandbox")) if sandbox_state is not None: sandbox_id = sandbox_state.get("sandbox_id") if sandbox_id is not None: diff --git a/backend/tests/test_ensure_sandbox_initialized.py b/backend/tests/test_ensure_sandbox_initialized.py new file mode 100644 index 000000000..572c6a52e --- /dev/null +++ b/backend/tests/test_ensure_sandbox_initialized.py @@ -0,0 +1,200 @@ +"""Tests for ensure_sandbox_initialized with fork-restored channel values.""" + +from __future__ import annotations + +import pytest +from langchain.tools import ToolRuntime +from langgraph.types import Overwrite + +from deerflow.sandbox.sandbox import Sandbox +from deerflow.sandbox.sandbox_provider import SandboxProvider, reset_sandbox_provider, set_sandbox_provider +from deerflow.sandbox.search import GrepMatch +from deerflow.sandbox.tools import ensure_sandbox_initialized, ensure_sandbox_initialized_async + + +class _StubSandbox(Sandbox): + def execute_command(self, command: str, env: dict | None = None, timeout: float | None = None) -> str: + del env, timeout + return "OK" + + def read_file(self, path: str) -> str: + return "content" + + def download_file(self, path: str) -> bytes: + return b"content" + + def list_dir(self, path: str, max_depth: int = 2) -> list[str]: + return ["/mnt/user-data/workspace/file.txt"] + + def write_file(self, path: str, content: str, append: bool = False) -> None: + return None + + def glob(self, path: str, pattern: str, *, include_dirs: bool = False, max_results: int = 200) -> tuple[list[str], bool]: + return [], False + + def grep( + self, + path: str, + pattern: str, + *, + glob: str | None = None, + literal: bool = False, + case_sensitive: bool = False, + max_results: int = 100, + ) -> tuple[list[GrepMatch], bool]: + return [], False + + def update_file(self, path: str, content: bytes) -> None: + return None + + +class _RecordingProvider(SandboxProvider): + def __init__(self) -> None: + self.sandbox = _StubSandbox("stub") + + def acquire(self, thread_id: str | None = None, *, user_id: str | None = None) -> str: + raise AssertionError("state already carries a sandbox; acquire must not run") + + async def acquire_async(self, thread_id: str | None = None, *, user_id: str | None = None) -> str: + raise AssertionError("state already carries a sandbox; acquire must not run") + + def get(self, sandbox_id: str) -> Sandbox | None: + if sandbox_id == "parent-sandbox": + return self.sandbox + return None + + def release(self, sandbox_id: str) -> None: + return None + + +class _FallthroughProvider(SandboxProvider): + """Provider whose parent id has expired, forcing a fresh acquire.""" + + def __init__(self) -> None: + self.sandbox = _StubSandbox("fresh") + self.acquired: list[str | None] = [] + + def acquire(self, thread_id: str | None = None, *, user_id: str | None = None) -> str: + self.acquired.append(thread_id) + return "fresh-sandbox" + + async def acquire_async(self, thread_id: str | None = None, *, user_id: str | None = None) -> str: + self.acquired.append(thread_id) + return "fresh-sandbox" + + def get(self, sandbox_id: str) -> Sandbox | None: + if sandbox_id == "fresh-sandbox": + return self.sandbox + return None + + def release(self, sandbox_id: str) -> None: + return None + + +def _make_runtime(state: dict) -> ToolRuntime: + return ToolRuntime( + state=state, + context={}, + config={"configurable": {}}, + stream_writer=lambda _: None, + tools=[], + tool_call_id="call-1", + store=None, + ) + + +def test_ensure_sandbox_initialized_unwraps_overwrite_state() -> None: + """Fork-restored state must not crash on the Overwrite wrapper.""" + provider = _RecordingProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({"sandbox": Overwrite({"sandbox_id": "parent-sandbox"})}) + sandbox = ensure_sandbox_initialized(runtime) + finally: + reset_sandbox_provider() + + assert sandbox is provider.sandbox + assert runtime.context["sandbox_id"] == "parent-sandbox" + # The reuse path must not take ownership: the wrapped state is left + # untouched, so after_agent still sees fork_restored and skips release. + assert isinstance(runtime.state["sandbox"], Overwrite) + + +@pytest.mark.anyio +async def test_ensure_sandbox_initialized_async_unwraps_overwrite_state() -> None: + provider = _RecordingProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({"sandbox": Overwrite({"sandbox_id": "parent-sandbox"})}) + sandbox = await ensure_sandbox_initialized_async(runtime) + finally: + reset_sandbox_provider() + + assert sandbox is provider.sandbox + assert runtime.context["sandbox_id"] == "parent-sandbox" + + +def test_ensure_sandbox_initialized_plain_state_unchanged() -> None: + provider = _RecordingProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({"sandbox": {"sandbox_id": "parent-sandbox"}}) + sandbox = ensure_sandbox_initialized(runtime) + finally: + reset_sandbox_provider() + + assert sandbox is provider.sandbox + assert runtime.context["sandbox_id"] == "parent-sandbox" + + +def test_ensure_sandbox_initialized_acquires_fresh_when_parent_missing() -> None: + """Acquire fall-through: the fork-restored id is gone from the provider, + so a fresh sandbox is acquired and the stale wrapped state is replaced + by the freshly acquired plain dict.""" + provider = _FallthroughProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({"sandbox": Overwrite({"sandbox_id": "parent-sandbox"})}) + runtime.context["thread_id"] = "t-1" + sandbox = ensure_sandbox_initialized(runtime) + finally: + reset_sandbox_provider() + + assert provider.acquired == ["t-1"] + assert sandbox is provider.sandbox + assert runtime.state["sandbox"] == {"sandbox_id": "fresh-sandbox"} + assert runtime.context["sandbox_id"] == "fresh-sandbox" + + +@pytest.mark.anyio +async def test_ensure_sandbox_initialized_async_plain_state_unchanged() -> None: + provider = _RecordingProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({"sandbox": {"sandbox_id": "parent-sandbox"}}) + sandbox = await ensure_sandbox_initialized_async(runtime) + finally: + reset_sandbox_provider() + + assert sandbox is provider.sandbox + assert runtime.context["sandbox_id"] == "parent-sandbox" + + +@pytest.mark.anyio +async def test_ensure_sandbox_initialized_async_acquires_fresh_when_parent_missing() -> None: + """Same fall-through as the sync path: the fork-restored id is gone from + the provider, so a fresh sandbox is acquired and the stale wrapped state + is replaced by the freshly acquired plain dict.""" + provider = _FallthroughProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({"sandbox": Overwrite({"sandbox_id": "parent-sandbox"})}) + runtime.context["thread_id"] = "t-1" + sandbox = await ensure_sandbox_initialized_async(runtime) + finally: + reset_sandbox_provider() + + assert provider.acquired == ["t-1"] + assert sandbox is provider.sandbox + assert runtime.state["sandbox"] == {"sandbox_id": "fresh-sandbox"} + assert runtime.context["sandbox_id"] == "fresh-sandbox"