fix(memory): drain persistent mutations on cancellation (#6092)

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>

backend/app/gateway/routers/memory.py

@@ -13,6 +13,7 @@ from deerflow.config.agents_config import AGENT_NAME_PATTERN
 from deerflow.config.memory_config import get_memory_config
 from deerflow.config.paths import make_safe_user_id
 from deerflow.runtime.user_context import get_effective_user_id
+from deerflow.utils.file_io import await_drained
 
 router = APIRouter(prefix="/api", tags=["memory"])
 
@@ -217,6 +218,11 @@ async def _get_memory_or_501(
         raise _map_memory_manager_error(exc) from exc
 
 
+async def _run_memory_mutation(func, /, *args, **kwargs):
+    """Run a persistent memory mutation without letting cancellation outlive it."""
+    return await await_drained(asyncio.to_thread(func, *args, **kwargs))
+
+
 class FactCreateRequest(BaseModel):
     """Request model for creating a memory fact."""
 
@@ -363,7 +369,7 @@ async def clear_memory(request: Request, agent_name: str | None = None) -> Memor
     selected_agent = _management_agent_name_or_501(manager, agent_name)
     scope_kwargs = _agent_scope_kwargs(selected_agent)
     try:
-        memory_data = await asyncio.to_thread(
+        memory_data = await _run_memory_mutation(
             manager.clear_memory,
             user_id=_resolve_memory_user_id(request),
             **scope_kwargs,
@@ -392,7 +398,7 @@ async def create_memory_fact_endpoint(body: FactCreateRequest, request: Request,
     selected_agent = _management_agent_name_or_501(manager, agent_name)
     scope_kwargs = _agent_scope_kwargs(selected_agent)
     try:
-        memory_data, fact_id = await asyncio.to_thread(
+        memory_data, fact_id = await _run_memory_mutation(
             manager.create_fact,
             content=body.content,
             category=body.category,
@@ -429,7 +435,7 @@ async def delete_memory_fact_endpoint(fact_id: str, request: Request, agent_name
     selected_agent = _management_agent_name_or_501(manager, agent_name)
     scope_kwargs = _agent_scope_kwargs(selected_agent)
     try:
-        memory_data = await asyncio.to_thread(
+        memory_data = await _run_memory_mutation(
             manager.delete_fact,
             fact_id,
             user_id=_resolve_memory_user_id(request),
@@ -461,7 +467,7 @@ async def update_memory_fact_endpoint(fact_id: str, body: FactPatchRequest, requ
     selected_agent = _management_agent_name_or_501(manager, agent_name)
     scope_kwargs = _agent_scope_kwargs(selected_agent)
     try:
-        memory_data = await asyncio.to_thread(
+        memory_data = await _run_memory_mutation(
             manager.update_fact,
             fact_id=fact_id,
             content=body.content,
@@ -519,7 +525,7 @@ async def import_memory(body: MemoryResponse, request: Request, agent_name: str
     selected_agent = _management_agent_name_or_501(manager, agent_name)
     scope_kwargs = _agent_scope_kwargs(selected_agent)
     try:
-        memory_data = await asyncio.to_thread(
+        memory_data = await _run_memory_mutation(
             manager.import_memory,
             body.model_dump(exclude_none=True),
             user_id=_resolve_memory_user_id(request),

backend/tests/test_memory_router.py

@@ -168,6 +168,26 @@ def test_import_memory_route_preserves_source_error() -> None:
     assert response.json()["facts"][0]["sourceError"] == "The agent previously suggested npm start."
 
 
+def test_clear_memory_routes_persistent_write_through_mutation_drain() -> None:
+    manager = MagicMock()
+    manager.clear_memory.return_value = _sample_memory()
+    calls: list[tuple] = []
+
+    async def drained(func, /, *args, **kwargs):
+        calls.append((func, args, kwargs))
+        return func(*args, **kwargs)
+
+    request = SimpleNamespace()
+    with (
+        patch("app.gateway.routers.memory.get_memory_manager", return_value=manager),
+        patch("app.gateway.routers.memory._resolve_memory_user_id", return_value="user-1"),
+        patch("app.gateway.routers.memory._run_memory_mutation", side_effect=drained),
+    ):
+        asyncio.run(call_unwrapped(memory.clear_memory, request))
+
+    assert calls == [(manager.clear_memory, (), {"user_id": "user-1"})]
+
+
 # ── clear ──────────────────────────────────────────────────────────────────