From 7eb41cfc93083267c8abd2384d315fb0594d32d7 Mon Sep 17 00:00:00 2001
From: yetuge <182123929+yetuge@users.noreply.github.com>
Date: Thu, 1 Oct 2026 20:13:57 +0800
Subject: [PATCH] refactor(gateway): hoist the persistent-write drain into one
 shared helper (#6151)

* refactor(gateway): hoist the persistent-write drain into one shared helper

Each persistent-write surface forked its own per-router drain wrapper:
_drained_write in routers/agents.py (action label, expected_errors
pass-through, lost-failure type-only logging), _run_store_mutation in
routers/subagents.py (bare drain), and the inline await_drained call in
routers/managed_models.py. Hoist the superset into
app.gateway.persistent_writes.run_drained_write and point all three
routers at it, per the review suggestion on #6087 (its family
predecessors #6078/#6092/#6093 have since landed).

The subagents and managed-models call sites now also emit the type-only
lost-failure log line on a drained-then-failed write (previously
silent), matching the tradeoff acknowledged in the #6087 review.
artifacts.py and personal_mcp.py keep their direct await_drained forms:
the former drains a coroutine with its own rollback flow, the latter
needs no label or expected-error handling.

Router tests move their monkeypatch targets to the shared name; the
lost-failure assertion pins the helper's logger. AGENTS.md points the
next persistent-write surface at the canonical implementation.

Signed-off-by: yetuge <2219677952@qq.com>

* fix(gateway): quiet expected domain errors and scope the AGENTS.md drain claim

Review follow-ups on #6151:

- The subagents create/update call sites pass (ManagedSubagentExistsError,)
  and (FileNotFoundError,) so the shared helper stays silent on routine
  4xx outcomes (duplicate-name 409, stale-store 404 race), matching
  create_agent's (AgentExistsError,).
- The managed-models save passes (HTTPException,), so the helper stops
  double-logging failures that _save_with_failure_logging already
  reports and stays silent on its routine 4xx outcomes.
- The AGENTS.md sentence now names the three converted routers instead
  of implying every persistent-write surface uses the helper; the
  shortened form also keeps the file under the guidance byte budget
  that agent-guidance enforces.

Signed-off-by: yetuge <2219677952@qq.com>

---------

Signed-off-by: yetuge <2219677952@qq.com>
---
 backend/app/gateway/AGENTS.md                 |  4 +-
 backend/app/gateway/persistent_writes.py      | 47 +++++++++++++++++++
 backend/app/gateway/routers/agents.py         | 43 ++---------------
 backend/app/gateway/routers/managed_models.py |  4 +-
 backend/app/gateway/routers/subagents.py      | 14 ++----
 .../tests/test_agents_router_write_drain.py   |  4 +-
 backend/tests/test_subagents_router.py        |  4 +-
 7 files changed, 64 insertions(+), 56 deletions(-)
 create mode 100644 backend/app/gateway/persistent_writes.py

diff --git a/backend/app/gateway/AGENTS.md b/backend/app/gateway/AGENTS.md
index eacbe30f67..d7a568b4bd 100644
--- a/backend/app/gateway/AGENTS.md
+++ b/backend/app/gateway/AGENTS.md
@@ -23,8 +23,8 @@ is drained across host cancellation before later runtime resources unwind; each
 service `stop()` remains bounded by its 30-second `asyncio.timeout`, so include
 that bound in pod grace-period budgeting.
 
-Managed-model saves drain on cancellation. Log worker failures before
-`await_drained` consumes them; omit exception text.
+Agents, subagents and managed-models routers drain persistent writes
+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
diff --git a/backend/app/gateway/persistent_writes.py b/backend/app/gateway/persistent_writes.py
new file mode 100644
index 0000000000..96bd1341fd
--- /dev/null
+++ b/backend/app/gateway/persistent_writes.py
@@ -0,0 +1,47 @@
+"""Shared helper for persistent router writes drained across cancellation.
+
+The persistent-write surfaces used to fork one drain wrapper each
+(``_drained_write`` in ``routers/agents.py``, ``_run_store_mutation`` in
+``routers/subagents.py``, the inline managed-models save drain). Route
+modules copy this one canonical implementation instead.
+"""
+
+import asyncio
+import logging
+from collections.abc import Callable
+
+from deerflow.utils.file_io import await_drained
+
+logger = logging.getLogger(__name__)
+
+
+async def run_drained_write[**P, T](
+    action: str,
+    func: Callable[P, T],
+    expected_errors: tuple[type[Exception], ...] = (),
+    /,
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> T:
+    """Run a persistent write off the event loop and drain it across cancellation.
+
+    A client that disconnects mid-request cancels the handler task: a bare
+    ``asyncio.to_thread`` either cancels a still-queued worker — the write
+    silently never happens — or detaches from a running one, dropping its
+    failure because the handler's ``except`` never runs. ``await_drained``
+    lets the worker finish first; expected domain errors re-raise unlogged
+    (the caller maps them to a 4xx while still connected), anything else is
+    logged with the exception type only — the text can carry user content.
+    """
+
+    def _logged() -> T:
+        try:
+            return func(*args, **kwargs)
+        except expected_errors:
+            raise
+        except Exception as exc:
+            # Non-cancelled failures are logged again by the outer route handler.
+            logger.error("%s failed (%s)", action, type(exc).__name__)
+            raise
+
+    return await await_drained(asyncio.to_thread(_logged))
diff --git a/backend/app/gateway/routers/agents.py b/backend/app/gateway/routers/agents.py
index 717a11ba6d..d73372cae8 100644
--- a/backend/app/gateway/routers/agents.py
+++ b/backend/app/gateway/routers/agents.py
@@ -3,7 +3,6 @@
 import asyncio
 import logging
 import re
-from collections.abc import Callable
 from pathlib import Path
 from typing import Literal
 
@@ -11,6 +10,7 @@ from fastapi import APIRouter, HTTPException, Request
 from pydantic import BaseModel, Field
 
 from app.gateway.authz import require_permission
+from app.gateway.persistent_writes import run_drained_write
 from deerflow.agents.memory.manager import get_memory_manager
 from deerflow.config.agents_api_config import get_agents_api_config
 from deerflow.config.agents_config import (
@@ -27,44 +27,11 @@ from deerflow.config.paths import get_paths
 from deerflow.knowledge_scope import KnowledgeScope, canonicalize_knowledge_scope
 from deerflow.persistence.agents import AgentDeleteOutcome, AgentExistsError, get_agent_store
 from deerflow.runtime.user_context import get_effective_user_id
-from deerflow.utils.file_io import await_drained
 
 logger = logging.getLogger(__name__)
 router = APIRouter(prefix="/api", tags=["agents"])
 
 
-async def _drained_write[**P, T](
-    action: str,
-    func: Callable[P, T],
-    expected_errors: tuple[type[Exception], ...] = (),
-    /,
-    *args: P.args,
-    **kwargs: P.kwargs,
-) -> T:
-    """Run a persistent write off the event loop and drain it across cancellation.
-
-    A client that disconnects mid-request cancels the handler task: a bare
-    ``asyncio.to_thread`` either cancels a still-queued worker — the write
-    silently never happens — or detaches from a running one, dropping its
-    failure because the handler's ``except`` never runs. ``await_drained``
-    lets the worker finish first; expected domain errors re-raise unlogged
-    (the caller maps them to a 4xx while still connected), anything else is
-    logged with the exception type only — the text can carry user content.
-    """
-
-    def _logged() -> T:
-        try:
-            return func(*args, **kwargs)
-        except expected_errors:
-            raise
-        except Exception as exc:
-            # Non-cancelled failures are logged again by the outer route handler.
-            logger.error("%s failed (%s)", action, type(exc).__name__)
-            raise
-
-    return await await_drained(asyncio.to_thread(_logged))
-
-
 AGENT_NAME_PATTERN = re.compile(r"^[A-Za-z0-9-]+$")
 
 ReasoningEffort = Literal["low", "medium", "high"]
@@ -411,7 +378,7 @@ async def create_agent_endpoint(body: AgentCreateRequest, request: Request) -> A
         return _agent_config_to_response(agent_cfg, include_soul=True, user_id=user_id)
 
     try:
-        return await _drained_write("Create agent", _create_agent, (AgentExistsError,))
+        return await run_drained_write("Create agent", _create_agent, (AgentExistsError,))
     except AgentExistsError:
         raise HTTPException(status_code=409, detail=f"Agent '{normalized_name}' already exists")
     except Exception as e:
@@ -536,7 +503,7 @@ async def update_agent(name: str, body: AgentUpdateRequest, request: Request) ->
             def _update_agent() -> None:
                 get_agent_store().update(name, updated, body.soul, user_id=user_id)
 
-            await _drained_write("Update agent", _update_agent)
+            await run_drained_write("Update agent", _update_agent)
 
         logger.info(f"Updated agent '{name}'")
 
@@ -635,7 +602,7 @@ async def update_user_profile(body: UserProfileUpdateRequest, request: Request)
         return user_md_path
 
     try:
-        user_md_path = await _drained_write("Update user profile", _write_profile)
+        user_md_path = await run_drained_write("Update user profile", _write_profile)
         logger.info(f"Updated USER.md at {user_md_path}")
         return UserProfileResponse(content=body.content or None)
     except Exception as e:
@@ -668,7 +635,7 @@ async def delete_agent(name: str, request: Request) -> None:
     try:
         # Off the event loop: resolve store + cancel → delete → cancel-on-success
         # (get_agent_store / memory manager do blocking config and FS I/O).
-        outcome = await _drained_write("Delete agent", _delete_agent_with_memory_cancel, (), name, user_id)
+        outcome = await run_drained_write("Delete agent", _delete_agent_with_memory_cancel, (), name, user_id)
     except Exception as e:
         logger.error(f"Failed to delete agent '{name}': {e}", exc_info=True)
         raise HTTPException(status_code=500, detail=f"Failed to delete agent: {str(e)}")
diff --git a/backend/app/gateway/routers/managed_models.py b/backend/app/gateway/routers/managed_models.py
index 19641edc65..4ca5eaa1cb 100644
--- a/backend/app/gateway/routers/managed_models.py
+++ b/backend/app/gateway/routers/managed_models.py
@@ -8,10 +8,10 @@ from langchain_core.messages import HumanMessage
 from pydantic import BaseModel, ConfigDict
 
 from app.gateway.deps import require_admin_user
+from app.gateway.persistent_writes import run_drained_write
 from deerflow.config.app_config import get_app_config
 from deerflow.config.managed_models import ManagedModel, ManagedModelStore
 from deerflow.reflection import resolve_class
-from deerflow.utils.file_io import await_drained
 
 logger = logging.getLogger(__name__)
 router = APIRouter(prefix="/api/managed-models", tags=["models"])
@@ -74,7 +74,7 @@ def _save_with_failure_logging(body: SaveModelRequest):
 @router.put("")
 async def save_model(request: Request, body: SaveModelRequest):
     await require_admin_user(request, detail=_ADMIN)
-    return await await_drained(asyncio.to_thread(_save_with_failure_logging, body))
+    return await run_drained_write("Save managed model", _save_with_failure_logging, (HTTPException,), body)
 
 
 def _probe_config(body: SaveModelRequest):
diff --git a/backend/app/gateway/routers/subagents.py b/backend/app/gateway/routers/subagents.py
index 421e5c5e77..948534cf90 100644
--- a/backend/app/gateway/routers/subagents.py
+++ b/backend/app/gateway/routers/subagents.py
@@ -3,13 +3,13 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Callable
 from typing import Any, Literal
 
 from fastapi import APIRouter, HTTPException, Request
 from pydantic import BaseModel, Field, ValidationError
 
 from app.gateway.deps import is_admin_user, require_admin_user
+from app.gateway.persistent_writes import run_drained_write
 from deerflow.config.app_config import get_app_config
 from deerflow.persistence.managed_subagents import (
     ManagedSubagentDefinition,
@@ -21,7 +21,6 @@ from deerflow.persistence.managed_subagents.base import (
     normalize_managed_subagent_name,
 )
 from deerflow.subagents.builtins import BUILTIN_SUBAGENTS
-from deerflow.utils.file_io import await_drained
 
 router = APIRouter(prefix="/api/subagents", tags=["subagents"])
 _ADMIN_REQUIRED_DETAIL = "Admin privileges are required to manage subagents."
@@ -76,11 +75,6 @@ class ManagedSubagentUpdateRequest(BaseModel):
     enabled: bool | None = None
 
 
-async def _run_store_mutation[**P, T](func: Callable[P, T], /, *args: P.args, **kwargs: P.kwargs) -> T:
-    """Drain persistent managed-subagent writes before cancellation propagates."""
-    return await await_drained(asyncio.to_thread(func, *args, **kwargs))
-
-
 def _validate_model(model: str, app_config) -> None:
     if model == "inherit":
         return
@@ -187,7 +181,7 @@ async def create_managed_subagent(request: Request, body: ManagedSubagentCreateR
         raise HTTPException(status_code=409, detail=f"Subagent name '{definition.name}' is reserved by a built-in or config.yaml definition.")
     store = get_managed_subagent_store(app_config)
     try:
-        await _run_store_mutation(store.create, definition)
+        await run_drained_write("Create managed subagent", store.create, (ManagedSubagentExistsError,), definition)
     except ManagedSubagentExistsError:
         raise HTTPException(status_code=409, detail=f"Managed subagent '{definition.name}' already exists")
     return SubagentResponse(
@@ -218,7 +212,7 @@ async def update_managed_subagent(name: str, request: Request, body: ManagedSuba
         raise HTTPException(status_code=422, detail=exc.errors()) from exc
     _validate_model(updated.model, app_config)
     try:
-        await _run_store_mutation(store.update, updated)
+        await run_drained_write("Update managed subagent", store.update, (FileNotFoundError,), updated)
     except FileNotFoundError:
         # The definition may be deleted by another administrator after the
         # read above. Preserve the endpoint's not-found contract instead of
@@ -241,5 +235,5 @@ async def delete_managed_subagent(name: str, request: Request) -> None:
     name = _validate_path_name(name)
     app_config = await asyncio.to_thread(get_app_config)
     store = get_managed_subagent_store(app_config)
-    if not await _run_store_mutation(store.delete, name):
+    if not await run_drained_write("Delete managed subagent", store.delete, (), name):
         raise HTTPException(status_code=404, detail=f"Managed subagent '{name}' not found")
diff --git a/backend/tests/test_agents_router_write_drain.py b/backend/tests/test_agents_router_write_drain.py
index 3f6346ab50..6e9f2d175b 100644
--- a/backend/tests/test_agents_router_write_drain.py
+++ b/backend/tests/test_agents_router_write_drain.py
@@ -62,7 +62,7 @@ async def test_agent_persistent_writes_route_through_write_drain(_agent_env, mon
         assert isinstance(expected_errors, tuple)
         return func(*args, **kwargs)
 
-    monkeypatch.setattr(router, "_drained_write", drained)
+    monkeypatch.setattr(router, "run_drained_write", drained)
 
     await create_agent_endpoint(AgentCreateRequest(name="planner", model="agent-model"))
     await update_agent("planner", AgentUpdateRequest(description="later"))
@@ -140,7 +140,7 @@ async def test_agent_update_logs_lost_worker_failure_after_cancellation(_agent_e
                 task.cancel()
                 await asyncio.gather(task, return_exceptions=True)
 
-    failures = [record for record in caplog.records if record.name == router.__name__ and "Update agent failed" in record.message]
+    failures = [record for record in caplog.records if record.name == "app.gateway.persistent_writes" and "Update agent failed" in record.message]
     assert len(failures) == 1
     assert "OSError" in failures[0].message
     assert "worker-secret" not in caplog.text
diff --git a/backend/tests/test_subagents_router.py b/backend/tests/test_subagents_router.py
index 9017d1fed9..5525e316da 100644
--- a/backend/tests/test_subagents_router.py
+++ b/backend/tests/test_subagents_router.py
@@ -59,12 +59,12 @@ async def test_managed_subagent_writes_route_through_mutation_drain(monkeypatch,
     store = Store()
     calls: list[tuple[object, tuple[object, ...]]] = []
 
-    async def drained(func, /, *args, **kwargs):
+    async def drained(action, func, expected_errors=(), /, *args, **kwargs):
         calls.append((func, args))
         return func(*args, **kwargs)
 
     monkeypatch.setattr(router, "get_managed_subagent_store", lambda *_: store)
-    monkeypatch.setattr(router, "_run_store_mutation", drained)
+    monkeypatch.setattr(router, "run_drained_write", drained)
 
     if operation == "create":
         await router.create_managed_subagent(