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 # ---------------------------------------------------------------------------