fix(gateway): close initialized sandbox provider on shutdown (#7347)

Co-authored-by: Totoro-qaq <279883115+Totoro-qaq@users.noreply.github.com>

README.md

@@ -513,6 +513,10 @@ single-label cluster hosts, and Docker/Podman internal hostnames do not inherit
 `HTTP_PROXY` or `HTTPS_PROXY`. External sandbox FQDNs and public IPs still
 honor environment proxy settings.
 
+After native runtime shutdown and MCP teardown, Gateway closes an initialized
+sandbox provider, including AIO sandboxes first created by a worker-thread tool call.
+It does not start an unused sandbox provider solely for cleanup.
+
 Backend processes automatically pick up `config.yaml` changes on the next config access, so model metadata updates do not require a manual restart during development.
 
 Gateway runs use the top-level `recursion_limit` in `config.yaml` when an API

backend/app/gateway/AGENTS.md

@@ -21,10 +21,10 @@ resources unwind; each service `stop()` has a 30-second `asyncio.timeout`.
 Routers with persistent writes drain through
 `app.gateway.persistent_writes.run_drained_write`.
 
-Gateway lifespan closes the MCP session pool after `langgraph_runtime` exits:
-its run drain and extension-service teardown can still make MCP calls during
-earlier shutdown hooks. Keep the pool close bounded and failure-isolated, but
-do not move it ahead of those session producers.
+After `langgraph_runtime` exits (run drain and extension teardown), close MCP
+bounded and failure-isolated. Then drain off-thread
+`shutdown_sandbox_provider()` in `finally`, including cancellation. Never
+initialize for cleanup; log failures and retain native retryability.
 
 Retention holds durable thread admission and the checkpoint lock before
 checkpoint reads. Inspect root head/parent and leaf writes only; protect selectors,

backend/app/gateway/app.py

@@ -584,29 +584,40 @@ async def _start_scheduled_task_notification_delivery(app: FastAPI, startup_conf
 
 @asynccontextmanager
 async def _runtime_with_mcp_pool_shutdown(app: FastAPI, startup_config: AppConfig) -> AsyncGenerator[None, None]:
-    """Close pooled MCP transports after runtime producers have stopped."""
+    """Close MCP transports and initialized sandboxes after runtime shutdown."""
     try:
         async with langgraph_runtime(app, startup_config):
             yield
     finally:
-        # RunManager drains active graph tasks when langgraph_runtime exits.
-        # Those tasks can still acquire new MCP sessions during earlier
-        # shutdown hooks, so closing the pool inside the runtime leaves fresh
-        # owner tasks and transports alive after the only close pass.
         try:
-            from deerflow.mcp.session_pool import get_session_pool
+            # RunManager drains active graph tasks when langgraph_runtime exits.
+            # Those tasks can still acquire new MCP sessions during earlier
+            # shutdown hooks, so closing the pool inside the runtime leaves fresh
+            # owner tasks and transports alive after the only close pass.
+            try:
+                from deerflow.mcp.session_pool import get_session_pool
 
-            await asyncio.wait_for(
-                get_session_pool().close_all(),
-                timeout=_SHUTDOWN_HOOK_TIMEOUT_SECONDS,
-            )
-        except TimeoutError:
-            logger.warning(
-                "MCP session pool shutdown exceeded %.1fs; proceeding with worker exit.",
-                _SHUTDOWN_HOOK_TIMEOUT_SECONDS,
-            )
-        except Exception:
-            logger.exception("Failed to close MCP sessions")
+                await asyncio.wait_for(
+                    get_session_pool().close_all(),
+                    timeout=_SHUTDOWN_HOOK_TIMEOUT_SECONDS,
+                )
+            except TimeoutError:
+                logger.warning(
+                    "MCP session pool shutdown exceeded %.1fs; proceeding with worker exit.",
+                    _SHUTDOWN_HOOK_TIMEOUT_SECONDS,
+                )
+            except Exception:
+                logger.exception("Failed to close MCP sessions")
+        finally:
+            # Lazy tool initialization may run off the main thread, where a
+            # provider cannot install signal handlers. Close only the existing
+            # provider after runtime shutdown and MCP teardown, without allocating.
+            try:
+                from deerflow.sandbox.sandbox_provider import shutdown_sandbox_provider
+
+                await await_drained(asyncio.to_thread(shutdown_sandbox_provider))
+            except Exception:
+                logger.exception("Failed to shut down sandbox provider")
 
 
 @asynccontextmanager

backend/tests/test_gateway_sandbox_shutdown.py

@@ -0,0 +1,197 @@
+"""Gateway owns sandbox cleanup even when lazy creation cannot bind signals."""
+
+import asyncio
+import threading
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+from fastapi import FastAPI
+from langchain_core.messages import ToolMessage
+from langgraph.prebuilt.tool_node import ToolCallRequest
+
+
+def _runtime(monkeypatch, events):
+    @asynccontextmanager
+    async def runtime(_app, _config):
+        try:
+            yield
+        finally:
+            events.append("runtime_drained")
+
+    monkeypatch.setattr("app.gateway.app.langgraph_runtime", runtime)
+    pool = SimpleNamespace(close_all=AsyncMock(side_effect=lambda: events.append("mcp_closed")))
+    monkeypatch.setattr("deerflow.mcp.session_pool.get_session_pool", lambda: pool)
+    return pool
+
+
+def test_shutdown_closes_provider_created_by_lazy_first_write_after_run_drain(monkeypatch):
+    from app.gateway.app import _runtime_with_mcp_pool_shutdown
+    from deerflow.agents.middlewares.read_before_write_middleware import ReadBeforeWriteMiddleware
+    from deerflow.runtime import RunManager, RunStatus
+    from deerflow.sandbox import sandbox_provider as registry
+
+    async def scenario():
+        events = []
+        loop_thread = threading.get_ident()
+
+        class Provider:
+            def __init__(self):
+                self.created_thread = threading.get_ident()
+                self.closed = False
+
+            def shutdown(self):
+                assert threading.get_ident() != loop_thread
+                events.append("sandbox_closed")
+                self.closed = True
+
+        monkeypatch.setattr(registry, "_default_sandbox_provider", None)
+        monkeypatch.setattr(registry, "get_app_config", lambda: SimpleNamespace(sandbox=SimpleNamespace(use="fixture:Provider")))
+        monkeypatch.setattr(registry, "resolve_class", lambda *_args: Provider)
+        pool = _runtime(monkeypatch, events)
+        manager = RunManager()
+        record = await manager.create("lazy-first-write")
+        await manager.set_status(record.run_id, RunStatus.running)
+        resume = asyncio.Event()
+        first_write_finished = asyncio.Event()
+        created = []
+
+        def read_new_file(_runtime, _path):
+            created.append(registry.get_sandbox_provider())
+            raise FileNotFoundError("new output")
+
+        middleware = ReadBeforeWriteMiddleware(content_reader=read_new_file)
+        state = {"messages": []}
+        runtime = SimpleNamespace(state=state, context={"thread_id": "lazy-first-write"})
+        call = ToolCallRequest(tool_call={"name": "write_file", "args": {"path": "/mnt/user-data/outputs/new.txt", "content": "new"}, "id": "write"}, tool=None, state=state, runtime=runtime)
+
+        async def write(_request):
+            return ToolMessage(content="OK", tool_call_id="write", name="write_file")
+
+        async def active_run():
+            await resume.wait()
+            await middleware.awrap_tool_call(call, write)
+            first_write_finished.set()
+            await asyncio.Event().wait()
+
+        @asynccontextmanager
+        async def draining_runtime(_app, _config):
+            record.task = asyncio.create_task(active_run())
+            try:
+                yield
+            finally:
+                resume.set()
+                await asyncio.wait_for(first_write_finished.wait(), timeout=1)
+                await manager.shutdown(timeout=1.0)
+                events.append("runtime_drained")
+
+        monkeypatch.setattr("app.gateway.app.langgraph_runtime", draining_runtime)
+        try:
+            async with _runtime_with_mcp_pool_shutdown(FastAPI(), SimpleNamespace()):
+                assert registry.get_initialized_sandbox_provider() is None
+            assert record.task.done()
+            assert len(created) == 1 and created[0].created_thread != loop_thread
+            assert created[0].closed
+            assert registry.get_initialized_sandbox_provider() is None
+            assert events == ["runtime_drained", "mcp_closed", "sandbox_closed"]
+            pool.close_all.assert_awaited_once()
+        finally:
+            resume.set()
+            if record.task is not None and not record.task.done():
+                record.task.cancel()
+                await asyncio.gather(record.task, return_exceptions=True)
+
+    asyncio.run(scenario())
+
+
+def test_shutdown_without_provider_does_not_allocate_one(monkeypatch):
+    from app.gateway.app import _runtime_with_mcp_pool_shutdown
+    from deerflow.sandbox import sandbox_provider as registry
+
+    async def scenario():
+        events = []
+        monkeypatch.setattr(registry, "_default_sandbox_provider", None)
+        monkeypatch.setattr(registry, "resolve_class", lambda *_args: pytest.fail("shutdown allocated a sandbox provider"))
+        _runtime(monkeypatch, events)
+        async with _runtime_with_mcp_pool_shutdown(FastAPI(), SimpleNamespace()):
+            assert registry.get_initialized_sandbox_provider() is None
+        assert registry.get_initialized_sandbox_provider() is None
+        assert events == ["runtime_drained", "mcp_closed"]
+
+    asyncio.run(scenario())
+
+
+def test_sandbox_shutdown_failure_is_logged_and_remains_retryable(monkeypatch, caplog):
+    from app.gateway.app import _runtime_with_mcp_pool_shutdown
+    from deerflow.sandbox import sandbox_provider as registry
+
+    async def scenario():
+        events = []
+        pool = _runtime(monkeypatch, events)
+
+        def fail_close():
+            raise RuntimeError("sandbox close failed")
+
+        provider = SimpleNamespace(shutdown=fail_close)
+        monkeypatch.setattr(registry, "_default_sandbox_provider", provider)
+        async with _runtime_with_mcp_pool_shutdown(FastAPI(), SimpleNamespace()):
+            pass
+        assert registry.get_initialized_sandbox_provider() is provider
+        assert "Failed to shut down sandbox provider" in caplog.text
+        pool.close_all.assert_awaited_once()
+
+    asyncio.run(scenario())
+
+
+@pytest.mark.parametrize("cancel_during", ["mcp", "sandbox"])
+def test_shutdown_drains_sandbox_worker_across_cancellation(monkeypatch, cancel_during):
+    from app.gateway.app import _runtime_with_mcp_pool_shutdown
+    from deerflow.sandbox import sandbox_provider as registry
+
+    async def scenario():
+        events = []
+        _runtime(monkeypatch, events)
+        pool_started = asyncio.Event()
+        worker_started = threading.Event()
+        release_worker = threading.Event()
+        worker_finished = threading.Event()
+        loop_thread = threading.get_ident()
+
+        async def close_pool():
+            events.append("mcp_close_started")
+            pool_started.set()
+            if cancel_during == "mcp":
+                await asyncio.Event().wait()
+
+        def close_sandbox():
+            assert threading.get_ident() != loop_thread
+            worker_started.set()
+            assert release_worker.wait(5.0)
+            worker_finished.set()
+
+        monkeypatch.setattr("deerflow.mcp.session_pool.get_session_pool", lambda: SimpleNamespace(close_all=close_pool))
+        monkeypatch.setattr(registry, "_default_sandbox_provider", SimpleNamespace(shutdown=close_sandbox))
+        context = _runtime_with_mcp_pool_shutdown(FastAPI(), SimpleNamespace())
+        await context.__aenter__()
+        task = asyncio.create_task(context.__aexit__(None, None, None))
+        try:
+            await asyncio.wait_for(pool_started.wait(), timeout=1)
+            if cancel_during == "mcp":
+                task.cancel()
+            assert await asyncio.to_thread(worker_started.wait, 1.0)
+            task.cancel()
+            await asyncio.sleep(0)
+            task.cancel()
+            await asyncio.sleep(0)
+            assert not task.done(), "shutdown abandoned the sandbox worker"
+            release_worker.set()
+            with pytest.raises(asyncio.CancelledError):
+                await task
+            assert worker_finished.is_set()
+            assert registry.get_initialized_sandbox_provider() is None
+        finally:
+            release_worker.set()
+            await asyncio.gather(task, return_exceptions=True)
+
+    asyncio.run(scenario())