fix(sandbox): unwrap Overwrite-wrapped state in ensure_sandbox_initialized (#4429)

* fix(sandbox): unwrap Overwrite-wrapped state in ensure_sandbox_initialized

The same fork-restored wrapper that crashed after_agent also reaches the sandbox init path, where sandbox_state.get() on the Overwrite object raises AttributeError. Share the unwrap helper from #4381's follow-up module deerflow/sandbox/overwrite.py and apply it at both init sites.

Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com>

* fix(sandbox): note why discarding fork_restored at the reuse sites is safe

* fix(sandbox): unify the Overwrite unwrap helper and pin the fall-through

- middleware.py now imports unwrap_sandbox from overwrite.py instead of
  keeping a second local copy whose docstring had already drifted; the
  shared helper covers both crash forms (subscript TypeError and the
  .get()-form AttributeError)
- test the acquire fall-through: when the fork-restored id is gone from
  the provider, a fresh sandbox is acquired and the stale wrapped state
  is replaced by the plain acquired dict
- the reuse-path test now also asserts runtime.state["sandbox"] stays
  wrapped, pinning the don't-treat-as-owned contract after_agent relies on

Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com>

* fix(sandbox): unwrap Overwrite state in the sibling sandbox readers

Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com>

---------

Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Co-authored-by: Willem Jiang <willem.jiang@gmail.com>

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:

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

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:

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"