From 6e6c078595e24579a6523a42b1cf4014245a92cf Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Sun, 26 Jul 2026 09:02:40 +0800 Subject: [PATCH] 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 --- .../harness/deerflow/sandbox/middleware.py | 34 ++++++++++- backend/tests/test_sandbox_middleware.py | 58 ++++++++++++++++++- 2 files changed, 88 insertions(+), 4 deletions(-) diff --git a/backend/packages/harness/deerflow/sandbox/middleware.py b/backend/packages/harness/deerflow/sandbox/middleware.py index a6ebbf7ea..9937bcdb6 100644 --- a/backend/packages/harness/deerflow/sandbox/middleware.py +++ b/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 diff --git a/backend/tests/test_sandbox_middleware.py b/backend/tests/test_sandbox_middleware.py index 53f61a080..ad8818add 100644 --- a/backend/tests/test_sandbox_middleware.py +++ b/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 # ---------------------------------------------------------------------------