diff --git a/README.md b/README.md index 21ce86899c..51bac9e8c1 100644 --- a/README.md +++ b/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 diff --git a/backend/app/gateway/AGENTS.md b/backend/app/gateway/AGENTS.md index 2da931f5ac..e05751a237 100644 --- a/backend/app/gateway/AGENTS.md +++ b/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, diff --git a/backend/app/gateway/app.py b/backend/app/gateway/app.py index 44c8381be9..5b8ac282cb 100644 --- a/backend/app/gateway/app.py +++ b/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 diff --git a/backend/tests/test_gateway_sandbox_shutdown.py b/backend/tests/test_gateway_sandbox_shutdown.py new file mode 100644 index 0000000000..1686ef8b71 --- /dev/null +++ b/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())