fix(gateway): keep run drain alive across repeated cancellation (#5487)

* test(gateway): cover repeated cancellation during run drain

* fix(gateway): keep run drain alive across repeated cancellation

* test(gateway): harden run-drain cancellation coverage

backend/app/gateway/deps.py

@@ -167,24 +167,37 @@ async def _drain_inflight_runs(run_manager: RunManager) -> None:
     """Drain in-flight runs before the checkpointer is torn down (issue #3373).
 
     Shields the (internally-bounded) drain so that even if the lifespan
-    coroutine is itself cancelled mid-shutdown — a second SIGINT or the server's
+    coroutine is repeatedly cancelled mid-shutdown — e.g. signal escalation or
-    graceful-shutdown timeout, i.e. the same signal storm behind #3373 — the
+    the server's graceful-shutdown timeout — the checkpointer pool is not closed
-    checkpointer pool is not closed while run tasks are still writing
+    while run tasks are still writing checkpoints. Cancellation is remembered
-    checkpoints. On such a cancellation we let the already-running drain finish
+    and propagated only after the already-running drain reaches a safe terminal
-    (it is bounded by ``RunManager.shutdown``'s own timeout) and then propagate
+    point.
-    the cancellation.
     """
     drain = asyncio.create_task(run_manager.shutdown(timeout=_RUN_DRAIN_TIMEOUT_SECONDS))
-    try:
+    cancellation: asyncio.CancelledError | None = None
-        await asyncio.shield(drain)
+
-    except asyncio.CancelledError:
+    while not drain.done():
-        # Re-shield so this second wait does not abandon the in-flight drain;
-        # it is bounded, so this cannot hang. Then re-raise to honour shutdown.
         try:
             await asyncio.shield(drain)
+        except asyncio.CancelledError as exc:
+            if cancellation is None:
+                cancellation = exc
+        except Exception:
+            if cancellation is not None:
+                logger.exception("In-flight run drain failed after shutdown cancellation")
+                raise cancellation
+            logger.exception("Failed to drain in-flight runs during shutdown")
+            return
+
+    if cancellation is not None:
+        try:
+            drain.result()
         except Exception:
             logger.exception("In-flight run drain failed after shutdown cancellation")
-        raise
+        raise cancellation
+
+    try:
+        drain.result()
     except Exception:
         logger.exception("Failed to drain in-flight runs during shutdown")
 
@@ -736,8 +749,8 @@ def get_run_context(request: Request) -> RunContext:
     ``app_config`` field is resolved live so per-run fields (e.g.
     ``models[*].max_tokens``) follow ``config.yaml`` edits; the
     ``event_store`` / ``run_events_config`` pair stays frozen to the snapshot
-    captured in :func:`langgraph_runtime` so callers never see a store bound
+    captured in :func:`langgraph_runtime` so callers never see a store bound to
-    to one backend paired with a config pointing at another.
+    one backend paired with a config pointing at another.
     """
     return RunContext(
         checkpointer=get_checkpointer(request),

backend/tests/test_gateway_run_drain_repeated_cancellation.py

@@ -0,0 +1,96 @@
+"""Repeated-cancellation regressions for Gateway run draining on shutdown."""
+
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_run_drain_waits_for_shutdown_across_repeated_cancellation():
+    """A second cancellation cannot let checkpointer teardown outrun the drain."""
+    from app.gateway import deps
+
+    shutdown_started = asyncio.Event()
+    allow_shutdown_finish = asyncio.Event()
+    shutdown_finished = asyncio.Event()
+
+    class _RunManager:
+        async def shutdown(self, *, timeout: float) -> None:
+            assert timeout == deps._RUN_DRAIN_TIMEOUT_SECONDS
+            shutdown_started.set()
+            await allow_shutdown_finish.wait()
+            shutdown_finished.set()
+
+    drain_task = asyncio.create_task(deps._drain_inflight_runs(_RunManager()))
+    await asyncio.wait_for(shutdown_started.wait(), timeout=1.0)
+
+    drain_task.cancel("first shutdown cancellation")
+    # Let the helper observe and remember the first cancellation, then suspend
+    # again while the still-blocked RunManager.shutdown() task remains alive.
+    await asyncio.sleep(0)
+    assert not drain_task.done()
+
+    # A second SIGINT / graceful-shutdown cancellation arrives while the helper
+    # is already draining after the first one. It must keep owning that wait so
+    # the surrounding AsyncExitStack cannot close the checkpointer underneath
+    # still-running run tasks.
+    drain_task.cancel("second shutdown cancellation")
+    await asyncio.sleep(0)
+    escaped_before_shutdown_finished = drain_task.done()
+
+    allow_shutdown_finish.set()
+    with pytest.raises(asyncio.CancelledError) as exc_info:
+        await drain_task
+
+    assert not escaped_before_shutdown_finished
+    assert shutdown_finished.is_set()
+    assert exc_info.value.args == ("first shutdown cancellation",)
+
+
+@pytest.mark.asyncio
+async def test_run_drain_logs_failure_and_returns_without_cancellation(caplog):
+    """A drain failure alone keeps the helper's existing best-effort contract."""
+    from app.gateway import deps
+
+    class _RunManager:
+        async def shutdown(self, *, timeout: float) -> None:
+            assert timeout == deps._RUN_DRAIN_TIMEOUT_SECONDS
+            raise RuntimeError("drain failed")
+
+    await deps._drain_inflight_runs(_RunManager())
+
+    assert "Failed to drain in-flight runs during shutdown" in caplog.text
+    assert "drain failed" in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_run_drain_failure_preserves_remembered_cancellation(caplog):
+    """A later drain error cannot replace cancellation already requested."""
+    from app.gateway import deps
+
+    shutdown_started = asyncio.Event()
+    allow_shutdown_failure = asyncio.Event()
+
+    class _RunManager:
+        async def shutdown(self, *, timeout: float) -> None:
+            assert timeout == deps._RUN_DRAIN_TIMEOUT_SECONDS
+            shutdown_started.set()
+            await allow_shutdown_failure.wait()
+            raise RuntimeError("drain failed after cancellation")
+
+    drain_task = asyncio.create_task(deps._drain_inflight_runs(_RunManager()))
+    await asyncio.wait_for(shutdown_started.wait(), timeout=1.0)
+
+    drain_task.cancel("shutdown requested")
+    await asyncio.sleep(0)
+    assert not drain_task.done()
+
+    allow_shutdown_failure.set()
+    with pytest.raises(asyncio.CancelledError) as exc_info:
+        await drain_task
+
+    assert exc_info.value.args == ("shutdown requested",)
+    assert "In-flight run drain failed after shutdown cancellation" in caplog.text
+    assert "drain failed after cancellation" in caplog.text