fix(gateway): drain managed model saves on cancellation (#6024)

* fix(gateway): drain managed model saves on cancellation

* test(gateway): pin managed model save cancellation drain

* fix(gateway): log failed drained managed model saves

Signed-off-by: Willem Jiang <willem.jiang@gmail.com>

* fix(gateway): keep guidance within hard size limit

---------

Signed-off-by: Willem Jiang <willem.jiang@gmail.com>
Co-authored-by: Willem Jiang <willem.jiang@gmail.com>

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

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):

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