fix(gateway): drain agent and user-profile writes across cancellation (#6087)

* fix(gateway): drain agent and user-profile writes across cancellation

A client that disconnects mid-request cancels the handler task. On the four
persistent-write endpoints of /api/agents - create, update, delete, and the
USER.md write - a bare asyncio.to_thread then either cancels a still-queued
worker (the write silently never happens) or detaches from a running one and
drops its failure, because the handler's except never runs.

Route the four writes through await_drained like the managed-subagent (#6023),
managed-model (#6024), and custom-skill (#6078) mutations, logging a lost
worker failure with the exception type only before the drain consumes it.
Expected domain errors (AgentExistsError on create) re-raise unlogged; the
caller still maps them to a 409 while connected. Reads stay bare: abandoning
them loses nothing.

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

* fix(gateway): address review nits on the agents write drain

Make expected_errors a positional-only parameter before *args so the
ParamSpec signature is checker-clean under PEP 612 (the one call site with
extra worker arguments passes () explicitly), acknowledge that
non-cancelled failures log twice like artifacts.py's drained commit does,
and drop the unused request stub helper from the tests.

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

---------

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

backend/app/gateway/routers/agents.py

@@ -3,6 +3,7 @@
 import asyncio
 import logging
 import re
+from collections.abc import Callable
 from pathlib import Path
 from typing import Literal
 
@@ -26,10 +27,44 @@ 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"]
@@ -376,7 +411,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 asyncio.to_thread(_create_agent)
+        return await _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:
@@ -501,7 +536,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 asyncio.to_thread(_update_agent)
+            await _drained_write("Update agent", _update_agent)
 
         logger.info(f"Updated agent '{name}'")
 
@@ -600,7 +635,7 @@ async def update_user_profile(body: UserProfileUpdateRequest, request: Request)
         return user_md_path
 
     try:
-        user_md_path = await asyncio.to_thread(_write_profile)
+        user_md_path = await _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:
@@ -633,7 +668,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 asyncio.to_thread(_delete_agent_with_memory_cancel, name, user_id)
+        outcome = await _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)}")

backend/tests/test_agents_router_write_drain.py

@@ -0,0 +1,189 @@
+"""Agent-router persistent writes drain across handler cancellation.
+
+The gateway convention (managed subagents, managed models, custom skills) is
+that a client disconnect must not cancel a queued persistence worker or
+silently drop its failure. The agent store writes (create / update / delete)
+and the USER.md write route through ``await_drained`` the same way; reads stay
+bare because abandoning them loses nothing.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import logging
+import threading
+from pathlib import Path
+
+import pytest
+
+from app.gateway.routers.agents import (
+    AgentCreateRequest,
+    AgentUpdateRequest,
+    UserProfileUpdateRequest,
+    create_agent_endpoint,
+    get_agent,
+    update_agent,
+    update_user_profile,
+)
+from deerflow.config.agents_api_config import load_agents_api_config_from_dict
+from deerflow.config.app_config import AppConfig, reset_app_config, set_app_config
+from deerflow.config.model_config import ModelConfig
+from deerflow.config.sandbox_config import SandboxConfig
+
+pytestmark = pytest.mark.asyncio
+
+
+@pytest.fixture
+def _agent_env(tmp_path: Path, monkeypatch):
+    monkeypatch.setenv("DEER_FLOW_HOME", str(tmp_path))
+    monkeypatch.setattr("deerflow.config.paths._paths", None)
+    load_agents_api_config_from_dict({"enabled": True})
+    set_app_config(
+        AppConfig(
+            models=[ModelConfig(name="agent-model", display_name="Agent Model", description=None, use="langchain_openai:ChatOpenAI", model="agent-model")],
+            sandbox=SandboxConfig(use="deerflow.sandbox.local:LocalSandboxProvider"),
+        )
+    )
+    try:
+        yield
+    finally:
+        load_agents_api_config_from_dict({})
+        reset_app_config()
+
+
+async def test_agent_persistent_writes_route_through_write_drain(_agent_env, monkeypatch):
+    from app.gateway.routers import agents as router
+
+    calls: list[str] = []
+
+    async def drained(action, func, expected_errors=(), /, *args, **kwargs):
+        calls.append(action)
+        assert isinstance(expected_errors, tuple)
+        return func(*args, **kwargs)
+
+    monkeypatch.setattr(router, "_drained_write", drained)
+
+    await create_agent_endpoint(AgentCreateRequest(name="planner", model="agent-model"))
+    await update_agent("planner", AgentUpdateRequest(description="later"))
+    await get_agent("planner")
+    await update_user_profile(UserProfileUpdateRequest(content="prefs"))
+    await router.delete_agent("planner")
+
+    assert calls == ["Create agent", "Update agent", "Update user profile", "Delete agent"]
+
+
+async def test_agent_update_drains_started_store_write_across_repeated_cancellation(_agent_env, monkeypatch):
+    from app.gateway.routers import agents as router
+
+    await create_agent_endpoint(AgentCreateRequest(name="planner", model="agent-model"))
+
+    started = threading.Event()
+    release = threading.Event()
+
+    class BlockingStore:
+        def update(self, *_args, **_kwargs):
+            started.set()
+            assert release.wait(timeout=5)
+
+    monkeypatch.setattr(router, "get_agent_store", lambda: BlockingStore())
+
+    task = asyncio.create_task(update_agent("planner", AgentUpdateRequest(description="later")))
+    try:
+        assert await asyncio.to_thread(started.wait, 5)
+        task.cancel()
+        await asyncio.sleep(0.05)
+        task.cancel()
+        await asyncio.sleep(0.05)
+        assert not task.done()
+
+        release.set()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+    finally:
+        release.set()
+        if not task.done():
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+
+
+async def test_agent_update_logs_lost_worker_failure_after_cancellation(_agent_env, monkeypatch, caplog):
+    from app.gateway.routers import agents as router
+
+    await create_agent_endpoint(AgentCreateRequest(name="planner", model="agent-model"))
+
+    started = threading.Event()
+    release = threading.Event()
+
+    class FailingStore:
+        def update(self, *_args, **_kwargs):
+            started.set()
+            assert release.wait(timeout=5)
+            raise OSError("worker-secret must not appear in logs")
+
+    monkeypatch.setattr(router, "get_agent_store", lambda: FailingStore())
+
+    with caplog.at_level(logging.ERROR, logger=router.__name__):
+        task = asyncio.create_task(update_agent("planner", AgentUpdateRequest(description="later")))
+        try:
+            assert await asyncio.to_thread(started.wait, 5)
+            task.cancel()
+            await asyncio.sleep(0)
+            assert not task.done()
+
+            release.set()
+            with pytest.raises(asyncio.CancelledError):
+                await task
+        finally:
+            release.set()
+            if not task.done():
+                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]
+    assert len(failures) == 1
+    assert "OSError" in failures[0].message
+    assert "worker-secret" not in caplog.text
+
+
+async def test_queued_user_profile_write_runs_despite_cancellation(_agent_env, tmp_path):
+    """A write queued behind a busy executor must still run when the caller goes away.
+
+    Without the drain, cancelling the handler cancels the still-queued worker
+    and the profile silently never lands; the run of this test on unfixed code
+    is exactly that data loss.
+    """
+    blocker_started = threading.Event()
+    release_blocker = threading.Event()
+
+    def _block():
+        blocker_started.set()
+        assert release_blocker.wait(timeout=5)
+
+    loop = asyncio.get_running_loop()
+    executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+    loop.set_default_executor(executor)
+    try:
+        jam = asyncio.create_task(asyncio.to_thread(_block))
+        # Poll without touching the executor: it has a single worker, and the
+        # blocker occupies it — an executor-based wait would deadlock here.
+        for _ in range(100):
+            if blocker_started.is_set():
+                break
+            await asyncio.sleep(0.05)
+        assert blocker_started.is_set()
+
+        task = asyncio.create_task(update_user_profile(UserProfileUpdateRequest(content="saved-prefs")))
+        await asyncio.sleep(0.1)  # the write is now queued behind the blocker
+        task.cancel()
+        await asyncio.sleep(0.1)
+
+        release_blocker.set()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+        await asyncio.gather(jam, return_exceptions=True)
+    finally:
+        release_blocker.set()
+        executor.shutdown(wait=False)
+
+    assert list(tmp_path.rglob("USER.md")), "queued USER.md write was lost to cancellation"