diff --git a/backend/app/gateway/routers/agents.py b/backend/app/gateway/routers/agents.py index 6f18369f78..717a11ba6d 100644 --- a/backend/app/gateway/routers/agents.py +++ b/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)}") diff --git a/backend/tests/test_agents_router_write_drain.py b/backend/tests/test_agents_router_write_drain.py new file mode 100644 index 0000000000..3f6346ab50 --- /dev/null +++ b/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"