fix(sandbox): unwrap Overwrite-wrapped sandbox state in after_agent (#4381)

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

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, and
after_agent/aafter_agent then crash subscripting the wrapper
("TypeError: 'Overwrite' object is not subscriptable") on the next
sandbox tool run in the forked conversation. Unwrap before reading the
sandbox id, and pin both hooks against an Overwrite-wrapped state.

Refs #4380 (bug 1 of 2; the history-loss half is a separate display path)

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

* fix(sandbox): don't release fork-restored parent sandboxes

The Overwrite-wrapped value replays the parent thread's sandbox state, so releasing it from the forked run would evict the parent's warm sandbox. _unwrap_sandbox now reports the wrapped form, and both after_agent hooks skip the release for it while keeping the normal path unchanged.

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,7 +9,7 @@ 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
+from langgraph.types import Command, Overwrite
 
 from deerflow.agents.thread_state import SandboxStateField, ThreadDataState
 from deerflow.runtime.user_context import resolve_runtime_user_id
@@ -25,6 +25,24 @@ 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.
 
@@ -99,9 +117,14 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
 
     @override
     def after_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None:
-        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:
+                # The wrapped value replays the parent thread's sandbox state;
+                # releasing it here would evict the parent's warm sandbox.
+                logger.info(f"Not releasing fork-restored sandbox {sandbox_id}")
+                return None
             logger.info(f"Releasing sandbox {sandbox_id}")
             get_sandbox_provider().release(sandbox_id)
             return None
@@ -117,9 +140,14 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
 
     @override
     async def aafter_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None:
-        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:
+                # The wrapped value replays the parent thread's sandbox state;
+                # releasing it here would evict the parent's warm sandbox.
+                logger.info(f"Not releasing fork-restored sandbox {sandbox_id}")
+                return None
             logger.info(f"Releasing sandbox {sandbox_id}")
             await self._release_sandbox_async(sandbox_id)
             return None

backend/tests/test_sandbox_middleware.py

@@ -9,7 +9,7 @@ from langchain.tools import ToolRuntime
 from langchain_core.messages import ToolMessage
 from langgraph.prebuilt.tool_node import ToolCallRequest
 from langgraph.runtime import Runtime
-from langgraph.types import Command
+from langgraph.types import Command, Overwrite
 
 from deerflow.agents.thread_state import ThreadState
 from deerflow.sandbox.middleware import SandboxMiddleware, SandboxMiddlewareState
@@ -252,6 +252,62 @@ async def test_aafter_agent_delegates_to_super_when_no_sandbox(monkeypatch: pyte
     assert calls == [(state, runtime)]
 
 
+def test_after_agent_unwraps_overwrite_sandbox_state() -> None:
+    """Fork-restored state may carry the sandbox channel Overwrite-wrapped."""
+    provider = _AsyncOnlyProvider()
+    set_sandbox_provider(provider)
+    try:
+        state = {"sandbox": Overwrite({"sandbox_id": "fork-restored"})}
+        result = SandboxMiddleware().after_agent(state, Runtime(context={}))
+    finally:
+        reset_sandbox_provider()
+
+    assert result is None
+    # The wrapped value replays the parent's sandbox; this run must not release it.
+    assert provider.released_ids == []
+
+
+def test_after_agent_releases_own_sandbox_state() -> None:
+    provider = _AsyncOnlyProvider()
+    set_sandbox_provider(provider)
+    try:
+        state = {"sandbox": {"sandbox_id": "own-sandbox"}}
+        result = SandboxMiddleware().after_agent(state, Runtime(context={}))
+    finally:
+        reset_sandbox_provider()
+
+    assert result is None
+    assert provider.released_ids == ["own-sandbox"]
+
+
+@pytest.mark.anyio
+async def test_aafter_agent_unwraps_overwrite_sandbox_state() -> None:
+    provider = _AsyncOnlyProvider()
+    set_sandbox_provider(provider)
+    try:
+        state = {"sandbox": Overwrite({"sandbox_id": "fork-restored"})}
+        result = await SandboxMiddleware().aafter_agent(state, Runtime(context={}))
+    finally:
+        reset_sandbox_provider()
+
+    assert result is None
+    assert provider.released_ids == []
+
+
+@pytest.mark.anyio
+async def test_aafter_agent_releases_own_sandbox_state() -> None:
+    provider = _AsyncOnlyProvider()
+    set_sandbox_provider(provider)
+    try:
+        state = {"sandbox": {"sandbox_id": "own-sandbox"}}
+        result = await SandboxMiddleware().aafter_agent(state, Runtime(context={}))
+    finally:
+        reset_sandbox_provider()
+
+    assert result is None
+    assert provider.released_ids == ["own-sandbox"]
+
+
 # ---------------------------------------------------------------------------
 # wrap_tool_call / awrap_tool_call: persistent sandbox state via Command
 # ---------------------------------------------------------------------------