diff --git a/backend/app/gateway/routers/subagents.py b/backend/app/gateway/routers/subagents.py index 35c0bb6436..421e5c5e77 100644 --- a/backend/app/gateway/routers/subagents.py +++ b/backend/app/gateway/routers/subagents.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +from collections.abc import Callable from typing import Any, Literal from fastapi import APIRouter, HTTPException, Request @@ -20,6 +21,7 @@ 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." @@ -74,6 +76,11 @@ 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 @@ -180,7 +187,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 asyncio.to_thread(store.create, definition) + await _run_store_mutation(store.create, definition) except ManagedSubagentExistsError: raise HTTPException(status_code=409, detail=f"Managed subagent '{definition.name}' already exists") return SubagentResponse( @@ -211,7 +218,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 asyncio.to_thread(store.update, updated) + await _run_store_mutation(store.update, updated) except FileNotFoundError: # The definition may be deleted by another administrator after the # read above. Preserve the endpoint's not-found contract instead of @@ -234,5 +241,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 asyncio.to_thread(store.delete, name): + if not await _run_store_mutation(store.delete, name): raise HTTPException(status_code=404, detail=f"Managed subagent '{name}' not found") diff --git a/backend/tests/test_subagents_router.py b/backend/tests/test_subagents_router.py index 40f48ea08f..9017d1fed9 100644 --- a/backend/tests/test_subagents_router.py +++ b/backend/tests/test_subagents_router.py @@ -2,6 +2,8 @@ from __future__ import annotations +import asyncio +import threading from types import SimpleNamespace import pytest @@ -31,6 +33,103 @@ def _environment(tmp_path, monkeypatch): reset_app_config() +@pytest.mark.parametrize("operation", ["create", "update", "delete"]) +async def test_managed_subagent_writes_route_through_mutation_drain(monkeypatch, operation): + definition = router.ManagedSubagentDefinition( + name="planner", + description="Plans work", + system_prompt="Plan the work.", + ) + + class Store: + def create(self, value): + return value + + def get(self, name): + assert name == "planner" + return definition + + def update(self, value): + return value + + def delete(self, name): + assert name == "planner" + return True + + store = Store() + calls: list[tuple[object, tuple[object, ...]]] = [] + + async def drained(func, /, *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) + + if operation == "create": + await router.create_managed_subagent( + _request("admin"), + router.ManagedSubagentCreateRequest( + name="planner", + description="Plans work", + system_prompt="Plan the work.", + ), + ) + expected = store.create + elif operation == "update": + await router.update_managed_subagent( + "planner", + _request("admin"), + router.ManagedSubagentUpdateRequest(enabled=False), + ) + expected = store.update + else: + await router.delete_managed_subagent("planner", _request("admin")) + expected = store.delete + + assert len(calls) == 1 + assert calls[0][0] == expected + + +async def test_create_managed_subagent_drains_started_write_across_repeated_cancellation(monkeypatch): + started = threading.Event() + release = threading.Event() + + class Store: + def create(self, _definition): + started.set() + assert release.wait(timeout=2) + + monkeypatch.setattr(router, "get_managed_subagent_store", lambda *_: Store()) + + task = asyncio.create_task( + router.create_managed_subagent( + _request("admin"), + router.ManagedSubagentCreateRequest( + name="planner", + description="Plans work", + system_prompt="Plan the work.", + ), + ) + ) + try: + assert await asyncio.to_thread(started.wait, 2) + 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_admin_can_create_update_and_delete_managed_subagent(): created = await router.create_managed_subagent( _request("admin"),