diff --git a/backend/app/gateway/routers/personal_mcp.py b/backend/app/gateway/routers/personal_mcp.py index 81585ea2de..f9e6d56b7a 100644 --- a/backend/app/gateway/routers/personal_mcp.py +++ b/backend/app/gateway/routers/personal_mcp.py @@ -12,6 +12,7 @@ from app.gateway.routers import mcp from deerflow.capabilities.runtime import ambiguous_installation_ids from deerflow.config.extensions_config import ExtensionsConfig, atomic_write_extensions_config, extensions_config_file_lock, extensions_config_write_lock from deerflow.mcp.user_config import read_user_mcp_config, user_mcp_config_path +from deerflow.utils.file_io import await_drained router = APIRouter(prefix="/api/mcp/personal/config", tags=["mcp"]) @@ -102,7 +103,7 @@ async def get_configuration(request: Request): async def _write(request: Request, operation: str, body): owner = await _owner(request) admin = await is_admin_user(request) - raw = await asyncio.to_thread(_mutate, owner, operation, body, admin=admin) + raw = await await_drained(asyncio.to_thread(_mutate, owner, operation, body, admin=admin)) return _response(raw) diff --git a/backend/tests/test_personal_mcp.py b/backend/tests/test_personal_mcp.py index 9e8c8aa367..e411b5eed4 100644 --- a/backend/tests/test_personal_mcp.py +++ b/backend/tests/test_personal_mcp.py @@ -48,6 +48,41 @@ def create(client, user, *, name="github", token=None, role="admin"): ) +@pytest.mark.asyncio +async def test_personal_config_write_drains_started_mutation_across_cancellation(monkeypatch): + started = asyncio.Event() + + monkeypatch.setattr(personal_mcp, "_owner", AsyncMock(return_value="alice")) + monkeypatch.setattr(personal_mcp, "is_admin_user", AsyncMock(return_value=True)) + + def mutate(*_args, **_kwargs): + started_loop.call_soon_threadsafe(started.set) + release_thread.wait(timeout=5) + return {"mcpServers": {}} + + import threading + + started_loop = asyncio.get_running_loop() + release_thread = threading.Event() + monkeypatch.setattr(personal_mcp, "_mutate", mutate) + + task = asyncio.create_task(personal_mcp._write(SimpleNamespace(), "delete", "missing")) + try: + await asyncio.wait_for(started.wait(), timeout=5) + task.cancel() + await asyncio.sleep(0) + assert not task.done() + + release_thread.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + release_thread.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + def test_persistent_same_name_connections_are_owner_only(personal_client): client = personal_client assert client.get("/api/mcp/personal/config").status_code == 401