From 1414e9fae75850480a8643d146d825b30b536714 Mon Sep 17 00:00:00 2001
From: NanPan <111261006+poijygfdyy@users.noreply.github.com>
Date: Thu, 1 Oct 2026 07:59:28 +0800
Subject: [PATCH] fix(memory): drain persistent mutations on cancellation
 (#6092)

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
---
 backend/app/gateway/routers/memory.py | 16 +++++++++++-----
 backend/tests/test_memory_router.py   | 20 ++++++++++++++++++++
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/backend/app/gateway/routers/memory.py b/backend/app/gateway/routers/memory.py
index 48454a7a7b..2d60ff7379 100644
--- a/backend/app/gateway/routers/memory.py
+++ b/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),
diff --git a/backend/tests/test_memory_router.py b/backend/tests/test_memory_router.py
index c09c3949fb..10ac04d883 100644
--- a/backend/tests/test_memory_router.py
+++ b/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 ──────────────────────────────────────────────────────────────────