From 11e0195fe9d8646c5fdbbe98138e086e64d7f6d1 Mon Sep 17 00:00:00 2001 From: yetuge <182123929+yetuge@users.noreply.github.com> Date: Fri, 2 Oct 2026 21:41:50 +0800 Subject: [PATCH] refactor(gateway): route memory mutations through the shared drained-write helper (#6164) Memory router kept its own bare drain wrapper (_run_memory_mutation, five call sites) after #6151 hoisted the superset into app.gateway.persistent_writes.run_drained_write - this is the remaining per-router fork called out in that review. Each call site passes the domain errors its handler already maps (NotImplementedError 501s, ValueError/KeyError 400/404/409 mappings, MemoryConflictError/MemoryCorruptionError, OSError 500s), so the helper stays silent on them exactly as the removed wrapper did; unexpected failures gain the type-only lost-failure log line per the AGENTS.md rule. Behavior is otherwise unchanged: same off-thread execution and cancel-drain semantics. The router test's drain stub moves to the shared name and pins the clear-memory error tuple. Signed-off-by: yetuge <2219677952@qq.com> --- backend/app/gateway/routers/memory.py | 27 ++++++++++++++++----------- backend/tests/test_memory_router.py | 10 ++++++---- 2 files changed, 22 insertions(+), 15 deletions(-) diff --git a/backend/app/gateway/routers/memory.py b/backend/app/gateway/routers/memory.py index 2d60ff7379..cb42d98fa2 100644 --- a/backend/app/gateway/routers/memory.py +++ b/backend/app/gateway/routers/memory.py @@ -8,12 +8,12 @@ from pydantic import BaseModel, Field, field_validator from app.gateway.authz import require_permission from app.gateway.internal_auth import get_trusted_internal_owner_user_id +from app.gateway.persistent_writes import run_drained_write from deerflow.agents.memory import MemoryConflictError, MemoryCorruptionError, MemoryManager, get_memory_manager 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"]) @@ -218,11 +218,6 @@ 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.""" @@ -369,8 +364,10 @@ 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 _run_memory_mutation( + memory_data = await run_drained_write( + "Clear memory data", manager.clear_memory, + (NotImplementedError, MemoryConflictError, MemoryCorruptionError, OSError), user_id=_resolve_memory_user_id(request), **scope_kwargs, ) @@ -398,8 +395,10 @@ 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 _run_memory_mutation( + memory_data, fact_id = await run_drained_write( + "Create memory fact", manager.create_fact, + (NotImplementedError, ValueError, MemoryConflictError, MemoryCorruptionError, OSError), content=body.content, category=body.category, confidence=body.confidence, @@ -435,8 +434,10 @@ 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 _run_memory_mutation( + memory_data = await run_drained_write( + "Delete memory fact", manager.delete_fact, + (NotImplementedError, KeyError, MemoryConflictError, MemoryCorruptionError, OSError), fact_id, user_id=_resolve_memory_user_id(request), **scope_kwargs, @@ -467,8 +468,10 @@ 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 _run_memory_mutation( + memory_data = await run_drained_write( + "Update memory fact", manager.update_fact, + (NotImplementedError, ValueError, KeyError, MemoryConflictError, MemoryCorruptionError, OSError), fact_id=fact_id, content=body.content, category=body.category, @@ -525,8 +528,10 @@ 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 _run_memory_mutation( + memory_data = await run_drained_write( + "Import memory data", manager.import_memory, + (NotImplementedError, ValueError, MemoryConflictError, MemoryCorruptionError, OSError), body.model_dump(exclude_none=True), user_id=_resolve_memory_user_id(request), **scope_kwargs, diff --git a/backend/tests/test_memory_router.py b/backend/tests/test_memory_router.py index 10ac04d883..ab5c63c058 100644 --- a/backend/tests/test_memory_router.py +++ b/backend/tests/test_memory_router.py @@ -173,19 +173,21 @@ def test_clear_memory_routes_persistent_write_through_mutation_drain() -> None: manager.clear_memory.return_value = _sample_memory() calls: list[tuple] = [] - async def drained(func, /, *args, **kwargs): - calls.append((func, args, kwargs)) + async def drained(action, func, expected_errors=(), /, *args, **kwargs): + calls.append((func, expected_errors, args, kwargs)) + assert isinstance(expected_errors, tuple) 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), + patch("app.gateway.routers.memory.run_drained_write", side_effect=drained), ): asyncio.run(call_unwrapped(memory.clear_memory, request)) - assert calls == [(manager.clear_memory, (), {"user_id": "user-1"})] + expected_errors = (NotImplementedError, MemoryConflictError, MemoryCorruptionError, OSError) + assert calls == [(manager.clear_memory, expected_errors, (), {"user_id": "user-1"})] # ── clear ──────────────────────────────────────────────────────────────────