diff --git a/backend/app/gateway/AGENTS.md b/backend/app/gateway/AGENTS.md index 6f165b4d9d..2a16b19702 100644 --- a/backend/app/gateway/AGENTS.md +++ b/backend/app/gateway/AGENTS.md @@ -23,6 +23,9 @@ 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. + Gateway lifespan closes the MCP session pool after `langgraph_runtime` exits: its run drain and extension-service teardown can still make MCP calls during earlier shutdown hooks. Keep the pool close bounded and failure-isolated, but diff --git a/backend/app/gateway/routers/managed_models.py b/backend/app/gateway/routers/managed_models.py index 5346f65d4c..19641edc65 100644 --- a/backend/app/gateway/routers/managed_models.py +++ b/backend/app/gateway/routers/managed_models.py @@ -1,6 +1,7 @@ """Admin-only shared model management. Credentials never leave the server.""" import asyncio +import logging from fastapi import APIRouter, HTTPException, Request from langchain_core.messages import HumanMessage @@ -10,7 +11,9 @@ from app.gateway.deps import require_admin_user 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"]) _ADMIN = "Admin privileges are required to manage shared models." @@ -54,10 +57,24 @@ def _save(body: SaveModelRequest): raise HTTPException(503, "Managed model storage is unavailable") from None +def _save_with_failure_logging(body: SaveModelRequest): + """Record save failures before a cancelled caller consumes the worker result.""" + try: + return _save(body) + except HTTPException as exc: + if exc.status_code >= 500: + logger.error("Managed model save failed (HTTP %d)", exc.status_code) + raise + except Exception as exc: + # Storage exceptions can contain credentials; retain only the type. + logger.error("Managed model save failed (%s)", type(exc).__name__) + raise + + @router.put("") async def save_model(request: Request, body: SaveModelRequest): await require_admin_user(request, detail=_ADMIN) - return await asyncio.to_thread(_save, body) + return await await_drained(asyncio.to_thread(_save_with_failure_logging, body)) def _probe_config(body: SaveModelRequest): diff --git a/backend/tests/test_managed_models.py b/backend/tests/test_managed_models.py index 519db0b3a8..0d5c91115b 100644 --- a/backend/tests/test_managed_models.py +++ b/backend/tests/test_managed_models.py @@ -1,5 +1,8 @@ """Managed models: persistence, snapshot resolution and administrator boundaries.""" +import asyncio +import logging +import threading from types import SimpleNamespace import pytest @@ -63,6 +66,83 @@ def test_endpoint_validation(url): ManagedModel(name="test", model="test", base_url=url) +@pytest.mark.asyncio +async def test_save_model_drains_started_persistence_across_cancellation(monkeypatch): + from app.gateway.routers import managed_models as router + + started = threading.Event() + release = threading.Event() + + def save(_body): + started.set() + assert release.wait(timeout=2) + return {"name": "managed-test", "has_api_key": True} + + monkeypatch.setattr(router, "_save", save) + request = SimpleNamespace(state=SimpleNamespace(user=SimpleNamespace(system_role="admin"))) + body = router.SaveModelRequest(config=profile()) + + task = asyncio.create_task(router.save_model(request, body)) + 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) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", [OSError, RuntimeError]) +async def test_save_model_logs_failed_drained_persistence_after_cancellation(store, monkeypatch, caplog, failure): + from app.gateway.routers import managed_models as router + + started = threading.Event() + release = threading.Event() + + def fail_save(self, *_args, **_kwargs): + started.set() + assert release.wait(timeout=5) + raise failure("test-secret must not appear in logs") + + monkeypatch.setattr(router, "get_app_config", lambda: AppConfig.model_validate({"sandbox": {"use": "test"}})) + monkeypatch.setattr(ManagedModelStore, "save", fail_save) + request = SimpleNamespace(state=SimpleNamespace(user=SimpleNamespace(system_role="admin"))) + body = router.SaveModelRequest(config=profile()) + + with caplog.at_level(logging.ERROR, logger=router.__name__): + task = asyncio.create_task(router.save_model(request, body)) + 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) + + save_errors = [record for record in caplog.records if record.name == router.__name__ and "Managed model save failed" in record.message] + assert len(save_errors) == 1 + expected_detail = "503" if failure is OSError else "RuntimeError" + assert expected_detail in save_errors[0].message + assert "test-secret" not in caplog.text + + @pytest.mark.asyncio async def test_admin_gate_and_response_redaction(store, monkeypatch): from app.gateway.routers import managed_models as router