[Router] Drain readiness before SIGTERM shutdown so k8s deregisters the pod first (#39016)

Co-authored-by: Kangyan Zhou <kangyan.zhou@radixark.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

experimental/sgl-router/tests/e2e/k8s_integration/conftest.py

@@ -109,40 +109,30 @@ def _wait_for_replacement_pod_ready(
     last_observed = "no pods"
 
     while time.time() < deadline:
-        result = _kubectl(
-            "get",
-            "pods",
-            "-n",
-            namespace,
-            "-l",
-            selector,
-            "-o",
-            "json",
-            check=False,
-        )
-        if getattr(result, "returncode", 0) == 0:
-            pods = json.loads(result.stdout or "{}").get("items", [])
-            names = [pod.get("metadata", {}).get("name", "") for pod in pods]
-            last_observed = ", ".join(filter(None, names)) or "no pods"
+        pods = _pods(selector, namespace, check=False)
+        names = [pod.get("metadata", {}).get("name", "") for pod in pods]
+        # An empty list is "nothing observed" whether the pods are gone or the
+        # kubectl call failed; both read the same in a timeout message.
+        last_observed = ", ".join(filter(None, names)) or "no pods"
 
-            if old_pod not in names:
-                for pod in sorted(
-                    pods, key=lambda item: item.get("metadata", {}).get("name", "")
+        if old_pod not in names:
+            for pod in sorted(
+                pods, key=lambda item: item.get("metadata", {}).get("name", "")
+            ):
+                metadata = pod.get("metadata", {})
+                status = pod.get("status", {})
+                ready = any(
+                    condition.get("type") == "Ready"
+                    and condition.get("status") == "True"
+                    for condition in status.get("conditions", [])
+                )
+                if (
+                    metadata.get("name") != old_pod
+                    and _is_live(pod)
+                    and status.get("phase") == "Running"
+                    and ready
                 ):
-                    metadata = pod.get("metadata", {})
-                    status = pod.get("status", {})
-                    ready = any(
-                        condition.get("type") == "Ready"
-                        and condition.get("status") == "True"
-                        for condition in status.get("conditions", [])
-                    )
-                    if (
-                        metadata.get("name") != old_pod
-                        and not metadata.get("deletionTimestamp")
-                        and status.get("phase") == "Running"
-                        and ready
-                    ):
-                        return metadata["name"]
+                    return metadata["name"]
 
         time.sleep(interval)
 
@@ -172,14 +162,20 @@ def _port_forward_start(
     service: str,
     local_port: int,
     remote_port: int,
+    resource: str = "svc",
 ) -> subprocess.Popen:
-    """Start kubectl port-forward and wait until the port is reachable."""
+    """Start kubectl port-forward and wait until the port is reachable.
+
+    `resource="pod"` binds one specific pod instead of the Service. A draining
+    pod is removed from the Service's ready endpoints, so a test that needs to
+    keep talking to it through the drain must address the pod directly.
+    """
     cmd = [
         "kubectl",
         "--context",
         KUBECTL_CONTEXT,
         "port-forward",
-        f"svc/{service}",
+        f"{resource}/{service}",
         f"{local_port}:{remote_port}",
         "-n",
         namespace,
@@ -215,6 +211,66 @@ def _cleanup_port_forward(name: str, pf: subprocess.Popen) -> None:
         logger.debug("Port-forward %s exited cleanly (rc=%s)", name, rc)
 
 
+def _pod_json(pod: str, namespace: str = NAMESPACE) -> dict:
+    """One pod's full object. The `or "{}"` mirrors
+    `_wait_for_replacement_pod_ready`: kubectl can hand back empty stdout, and a
+    JSONDecodeError there says nothing about what went wrong."""
+    result = _kubectl("get", "pod", pod, "-n", namespace, "-o", "json")
+    return json.loads(result.stdout or "{}")
+
+
+def _pods(
+    selector: str,
+    namespace: str = NAMESPACE,
+    check: bool = True,
+) -> list[dict]:
+    """Pod objects matching `selector`. `check=False` yields `[]` on a failed
+    kubectl instead of raising, for poll loops that expect the API server to be
+    briefly unavailable mid-rollout. The `or "{}"` guards kubectl handing back
+    empty stdout, where a JSONDecodeError would say nothing about what went
+    wrong."""
+    result = _kubectl(
+        "get", "pods", "-n", namespace, "-l", selector, "-o", "json", check=check
+    )
+    if getattr(result, "returncode", 0) != 0:
+        return []
+    return json.loads(result.stdout or "{}").get("items", [])
+
+
+def _is_live(pod: dict) -> bool:
+    """Whether a pod object is not already terminating. One predicate rather
+    than two copies of `deletionTimestamp`, so the replacement-pod poll and
+    `_pod_names` cannot drift apart on what counts as gone."""
+    return not pod.get("metadata", {}).get("deletionTimestamp")
+
+
+def _pod_names(selector: str, namespace: str = NAMESPACE) -> list[str]:
+    """Names of pods matching `selector`, excluding any already terminating."""
+    return [p["metadata"]["name"] for p in _pods(selector, namespace) if _is_live(p)]
+
+
+def _container_restart_count(
+    pod: str,
+    container: str,
+    namespace: str = NAMESPACE,
+) -> int:
+    """`restartCount` for one container — how a test observes that the process
+    exited and kubelet restarted it in place (no new pod, same name)."""
+    statuses = _pod_json(pod, namespace).get("status", {}).get("containerStatuses", [])
+    for status in statuses:
+        if status["name"] == container:
+            return int(status["restartCount"])
+    raise AssertionError(f"container {container!r} not found on pod {pod!r}")
+
+
+def _pod_ready_condition(pod: str, namespace: str = NAMESPACE) -> str:
+    """The pod's `Ready` condition as k8s currently sees it ("True"/"False")."""
+    for cond in _pod_json(pod, namespace).get("status", {}).get("conditions", []):
+        if cond["type"] == "Ready":
+            return cond["status"]
+    return "Unknown"
+
+
 def _poll_until(
     predicate,
     description: str,

experimental/sgl-router/tests/e2e/k8s_integration/manifests/router-cluster-scoped.yaml

@@ -16,6 +16,9 @@ spec:
         app: sgl-router-cluster
     spec:
       serviceAccountName: sgl-router-cluster
+      # Must exceed --shutdown-drain-secs plus the time in-flight requests need
+      # after the pause, or the pod is SIGKILLed mid-drain.
+      terminationGracePeriodSeconds: 40
       containers:
         - name: router
           image: sgl-router:e2e
@@ -37,6 +40,10 @@ spec:
             - "--service-discovery"
             - "--selector"
             - "app=sglang,cross-ns-test=true"
+            # Parity with router.yaml so the cross-namespace router exercises
+            # the same shutdown path; no test asserts on it here.
+            - "--shutdown-drain-secs"
+            - "8"
           ports:
             - containerPort: 8091
               name: http

experimental/sgl-router/tests/e2e/k8s_integration/manifests/router.yaml

@@ -14,6 +14,10 @@ spec:
         app: sgl-router
     spec:
       serviceAccountName: sgl-router
+      # Must exceed --shutdown-drain-secs plus the time in-flight requests
+      # need after the pause, or the pod is SIGKILLed mid-drain and the drain
+      # has bought nothing.
+      terminationGracePeriodSeconds: 40
       containers:
         - name: router
           image: sgl-router:e2e
@@ -44,6 +48,18 @@ spec:
             - "sgl-router-test"
             - "--selector"
             - "app=sglang"
+            # On SIGTERM, keep serving with /readyz at 503 so the endpoint
+            # removal reaches kube-proxy before the listener closes. Sized for
+            # the deletionTimestamp path, which does not wait on a probe.
+            # Probe-driven deregistration would instead need this above the
+            # readinessProbe's failureThreshold * periodSeconds below.
+            # test_shutdown_drain.py reads this value; do not restate it there.
+            - "--shutdown-drain-secs"
+            - "8"
+            # Declared so the startup advisory compares the drain against this
+            # pod's real grace period instead of assuming the k8s default.
+            - "--termination-grace-secs"
+            - "40"
           ports:
             - containerPort: 8090
               name: http

experimental/sgl-router/tests/e2e/k8s_integration/test_conftest_helpers.py

@@ -2,6 +2,7 @@ import json
 from types import SimpleNamespace
 
 import conftest as k8s_conftest
+import pytest
 
 
 def _pod(name: str, phase: str, ready: bool) -> dict:
@@ -52,3 +53,97 @@ def test_wait_for_replacement_pod_ignores_old_and_pending_pods(monkeypatch):
     assert replacement == new_pod
     assert len(calls) == 4
     assert all("-o" in args and "json" in args for args, _ in calls)
+
+
+def _stub_kubectl(monkeypatch, payload):
+    """Point every conftest helper at a canned kubectl response. `payload` is
+    the object kubectl would have printed."""
+    monkeypatch.setattr(
+        k8s_conftest,
+        "_kubectl",
+        lambda *args, **kwargs: SimpleNamespace(stdout=json.dumps(payload)),
+    )
+
+
+def test_pod_names_excludes_terminating_pods(monkeypatch):
+    terminating = _pod("sgl-router-old", "Running", True)
+    terminating["metadata"]["deletionTimestamp"] = "2026-01-01T00:00:00Z"
+    _stub_kubectl(
+        monkeypatch,
+        {"items": [terminating, _pod("sgl-router-new", "Running", True)]},
+    )
+
+    assert k8s_conftest._pod_names("app=sgl-router") == ["sgl-router-new"]
+
+
+def test_pods_returns_empty_on_a_failed_kubectl(monkeypatch):
+    """The poll loops call this with `check=False` precisely because the API
+    server can be briefly unavailable mid-rollout; a non-zero return must read
+    as "nothing observed", not raise out of the loop."""
+    monkeypatch.setattr(
+        k8s_conftest,
+        "_kubectl",
+        lambda *args, **kwargs: SimpleNamespace(stdout="", returncode=1),
+    )
+
+    assert k8s_conftest._pods("app=sgl-router", check=False) == []
+
+
+def test_pods_tolerates_empty_stdout(monkeypatch):
+    monkeypatch.setattr(
+        k8s_conftest,
+        "_kubectl",
+        lambda *args, **kwargs: SimpleNamespace(stdout=""),
+    )
+
+    assert k8s_conftest._pods("app=sgl-router") == []
+
+
+def test_container_restart_count_reads_the_named_container(monkeypatch):
+    _stub_kubectl(
+        monkeypatch,
+        {
+            "status": {
+                "containerStatuses": [
+                    {"name": "sidecar", "restartCount": 9},
+                    {"name": "router", "restartCount": 3},
+                ]
+            }
+        },
+    )
+
+    assert k8s_conftest._container_restart_count("sgl-router-0", "router") == 3
+
+
+def test_container_restart_count_rejects_a_missing_container(monkeypatch):
+    """`containerStatuses` lags during a restart, so the absent case is a real
+    state — and the drain test reads its whole timing floor off this number.
+    Returning 0 there would silently read as "never restarted"."""
+    _stub_kubectl(monkeypatch, {"status": {"containerStatuses": []}})
+
+    with pytest.raises(AssertionError, match="router"):
+        k8s_conftest._container_restart_count("sgl-router-0", "router")
+
+
+def test_pod_ready_condition_reports_the_ready_status(monkeypatch):
+    _stub_kubectl(
+        monkeypatch,
+        {
+            "status": {
+                "conditions": [
+                    {"type": "Initialized", "status": "True"},
+                    {"type": "Ready", "status": "False"},
+                ]
+            }
+        },
+    )
+
+    assert k8s_conftest._pod_ready_condition("sgl-router-0") == "False"
+
+
+def test_pod_ready_condition_is_unknown_before_the_condition_exists(monkeypatch):
+    """A pod whose Ready condition has not been written yet must read as
+    "Unknown" rather than crash the drain test's diagnostic logging."""
+    _stub_kubectl(monkeypatch, {"status": {"conditions": []}})
+
+    assert k8s_conftest._pod_ready_condition("sgl-router-0") == "Unknown"

experimental/sgl-router/tests/e2e/k8s_integration/test_shutdown_drain.py

@@ -0,0 +1,258 @@
+"""SIGTERM readiness-drain integration tests.
+
+The drain exists to produce a *Kubernetes* behaviour, and the Rust tests can
+only argue it: they substitute a channel for the real `Signal` and an
+in-process `AppContext` for a real pod. These run the shipped container, so
+they cover `main.rs::shutdown_signal` — the signal handler, the SIGTERM/SIGINT
+branch, and the `ctx` wiring — which no in-process test reaches.
+
+Why `kill -TERM 1` rather than `kubectl delete pod`: deleting a pod stamps a
+`deletionTimestamp`, and the endpoints controller marks the endpoint not-ready
+on that alone, without ever consulting `/readyz`. A delete-based test would
+therefore pass identically with the drain removed — it would look like
+coverage while pinning nothing. Signalling the process directly leaves the pod
+undeleted, so a `/readyz` 503 can only have come from the drain calling
+`AppContext::mark_not_ready`.
+
+The router image is `debian:bookworm-slim` with an exec-form ENTRYPOINT, so
+the binary is PID 1 and `kill -TERM 1` reaches it exactly as kubelet's SIGTERM
+would.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+import time
+from pathlib import Path
+
+import httpx
+from conftest import (
+    NAMESPACE,
+    _cleanup_port_forward,
+    _container_restart_count,
+    _kubectl,
+    _pod_names,
+    _pod_ready_condition,
+    _poll_until,
+    _port_forward_start,
+    _wait_for_deployment_ready,
+)
+
+logger = logging.getLogger(__name__)
+
+# Distinct from the shared 8090 forward so this test's pod-scoped forward
+# cannot collide with a leaked service-scoped one from another test.
+DRAIN_PORT = 8094
+
+_ROUTER_MANIFEST = Path(__file__).parent / "manifests" / "router.yaml"
+
+
+def _manifest_drain_secs() -> int:
+    """Read `--shutdown-drain-secs` out of the manifest the pod is started
+    from. Read rather than restated, because every assertion below is scaled to
+    the drain window: a manifest edit that this file did not track would leave
+    the test green while measuring the wrong window."""
+    args = _ROUTER_MANIFEST.read_text()
+    match = re.search(
+        r'"--shutdown-drain-secs"\s*\n\s*-\s*"(\d+)"',
+        args,
+    )
+    assert match, f"--shutdown-drain-secs not found in {_ROUTER_MANIFEST}"
+    return int(match.group(1))
+
+
+CONFIGURED_DRAIN_SECS = _manifest_drain_secs()
+
+
+def _manifest_grace_secs() -> tuple[int, int]:
+    """`terminationGracePeriodSeconds` from the pod spec, and the
+    `--termination-grace-secs` the router is told about it. Two places by
+    necessity — the router cannot read its own pod spec — which is exactly why
+    they can drift apart."""
+    manifest = _ROUTER_MANIFEST.read_text()
+    spec = re.search(r"terminationGracePeriodSeconds:\s*(\d+)", manifest)
+    assert spec, f"terminationGracePeriodSeconds not found in {_ROUTER_MANIFEST}"
+    declared = re.search(
+        r'"--termination-grace-secs"\s*\n\s*-\s*"(\d+)"',
+        manifest,
+    )
+    assert declared, f"--termination-grace-secs not found in {_ROUTER_MANIFEST}"
+    return int(spec.group(1)), int(declared.group(1))
+
+
+def test_declared_grace_period_matches_the_pod_spec():
+    """`--termination-grace-secs` silences the startup advisory, so a value
+    that has drifted from the pod's real `terminationGracePeriodSeconds` is
+    worse than no flag at all: it silences the warning against a budget the pod
+    does not have. No cluster needed — this is a manifest self-consistency
+    check, and it is the only thing standing between the two numbers."""
+    spec_secs, declared_secs = _manifest_grace_secs()
+    assert declared_secs == spec_secs, (
+        f"--termination-grace-secs is {declared_secs} but the pod spec grants "
+        f"{spec_secs}s; the advisory would be checked against the wrong budget"
+    )
+    assert CONFIGURED_DRAIN_SECS < spec_secs, (
+        f"the {CONFIGURED_DRAIN_SECS}s drain leaves no room under the {spec_secs}s "
+        f"grace period for the in-flight drain that follows it"
+    )
+
+
+# Budget for observing the /readyz flip, deliberately a fraction of the drain:
+# the assertions that follow it must still land inside the window, so the poll
+# cannot be allowed to consume the whole thing.
+FLIP_OBSERVATION_SECS = max(2, CONFIGURED_DRAIN_SECS // 2)
+
+# Floor on a mid-drain HTTP timeout. Below this the request has no realistic
+# chance on a loaded kind runner, so there is no point issuing it — the window
+# has effectively closed and `_mid_drain_timeout` says so instead.
+MIN_HTTP_TIMEOUT_SECS = 1.0
+
+# How long past the window the container restart may take to become VISIBLE.
+# Kubelet's own restart latency lands in here, and it only ever makes the
+# observed time longer — so this is slack on the measurement, not a second
+# claim about the drain. Sized to still catch a units regression that
+# LENGTHENS the pause: the `from_secs`/`from_millis` slip that
+# `ServerConfig::shutdown_drain()` exists to guard cuts both ways, and 8s
+# becoming 80s satisfies every lower bound in this file.
+RESTART_OBSERVATION_SLACK_SECS = 60
+
+
+def _mid_drain_timeout(sigterm_at: float, what: str, want: float) -> float:
+    """An HTTP timeout for a mid-drain assertion that cannot outlast the window
+    the assertion claims to run inside.
+
+    Without this the per-request timeouts sum past the drain (a 4s flip poll
+    plus 5s and 10s requests against an 8s window), so on a slow runner the
+    listener closes with a request still open and the test dies on whichever
+    transport error that raised — not on the assertion written to explain the
+    outcome. Checking the remaining budget up front puts the explanation back.
+    """
+    remaining = CONFIGURED_DRAIN_SECS - (time.monotonic() - sigterm_at)
+    assert remaining > MIN_HTTP_TIMEOUT_SECS, (
+        f"no drain window left for {what}: {CONFIGURED_DRAIN_SECS - remaining:.1f}s "
+        f"of the {CONFIGURED_DRAIN_SECS}s window already spent. If this runner is "
+        f"simply slow, raise --shutdown-drain-secs in {_ROUTER_MANIFEST.name}"
+    )
+    return min(want, remaining)
+
+
+def _router_pod() -> str:
+    pods = _pod_names("app=sgl-router")
+    assert len(pods) == 1, f"expected exactly one live router pod, got {pods}"
+    return pods[0]
+
+
+class TestReadinessDrain:
+    """SIGTERM must flip /readyz to 503 while the pod keeps serving."""
+
+    def test_sigterm_flips_readyz_while_the_pod_keeps_serving(self, k8s_cluster):
+        _wait_for_deployment_ready("sgl-router")
+        pod = _router_pod()
+        restarts_before = _container_restart_count(pod, "router")
+
+        # Bind the pod, not the Service: a draining pod leaves the Service's
+        # ready endpoints, and the point of this test is to keep talking to it
+        # after that happens.
+        pf = _port_forward_start(NAMESPACE, pod, DRAIN_PORT, 8090, resource="pod")
+        base = f"http://127.0.0.1:{DRAIN_PORT}"
+        try:
+            assert httpx.get(f"{base}/readyz", timeout=5.0).status_code == 200, (
+                "router must be ready before SIGTERM"
+            )
+
+            # Through `sh -c`: the slim image ships no `kill` binary, and
+            # `kubectl exec` execs directly rather than through a shell, so the
+            # builtin is the only way to signal PID 1 from outside.
+            sigterm_at = time.monotonic()
+            _kubectl("exec", "-n", NAMESPACE, pod, "--", "sh", "-c", "kill -TERM 1")
+
+            # The flip is observable from outside the pod.
+            _poll_until(
+                lambda: httpx.get(f"{base}/readyz", timeout=3.0).status_code == 503,
+                "/readyz returns 503 after SIGTERM",
+                timeout=FLIP_OBSERVATION_SECS,
+                interval=0.2,
+            )
+
+            # ...and the pod is still serving while it reports not-ready.
+            # `/healthz` staying 200 is what stops the liveness probe
+            # restarting a pod that is draining on purpose.
+            healthz_timeout = _mid_drain_timeout(sigterm_at, "the liveness probe", 5.0)
+            assert (
+                httpx.get(f"{base}/healthz", timeout=healthz_timeout).status_code == 200
+            ), "liveness must stay green while the pod drains"
+
+            # A proxied completion still succeeding does double duty: it is the
+            # request k8s may still route during the window, AND it proves the
+            # worker registry is non-empty — so the 503 above can only be the
+            # readiness flip, not `/readyz`'s other term.
+            chat = httpx.post(
+                f"{base}/v1/chat/completions",
+                json={
+                    "model": "tiny",
+                    "messages": [{"role": "user", "content": "drain"}],
+                },
+                timeout=_mid_drain_timeout(sigterm_at, "a proxied completion", 10.0),
+            )
+            assert chat.status_code == 200, (
+                f"a proxied request must still succeed mid-drain, got {chat.status_code}"
+            )
+
+            # Everything above claims to have run *inside* the window. Say so,
+            # so an overrun reads as "the window closed" and not as whichever
+            # transport error the closed listener happened to raise next.
+            mid_drain_elapsed = time.monotonic() - sigterm_at
+            assert mid_drain_elapsed < CONFIGURED_DRAIN_SECS, (
+                f"the mid-drain assertions took {mid_drain_elapsed:.1f}s, past the "
+                f"{CONFIGURED_DRAIN_SECS}s window they claim to observe"
+            )
+
+            # Recorded, not asserted: k8s needs failureThreshold consecutive
+            # failing probes, periodSeconds apart, to mark the pod not-ready —
+            # longer than the drain at the values in router.yaml. That is
+            # exactly why the default is sized for the deletionTimestamp path
+            # instead, and why probe-driven setups must raise it.
+            logger.info("pod Ready condition mid-drain: %s", _pod_ready_condition(pod))
+
+            # The drain's FLOOR, pinned where it is actually observable: hold
+            # until just inside the window and prove the process is still up.
+            # Timing the floor off the restart instead is satisfiable by test
+            # overhead alone — the restart poll does not start until everything
+            # above has run, so a build whose pause was 1s would still look like
+            # it lasted the whole window.
+            still_up_at = CONFIGURED_DRAIN_SECS - 1
+            time.sleep(max(0.0, still_up_at - (time.monotonic() - sigterm_at)))
+            assert _container_restart_count(pod, "router") == restarts_before, (
+                f"the router exited within {still_up_at}s of SIGTERM, short of the "
+                f"configured {CONFIGURED_DRAIN_SECS}s drain"
+            )
+
+        finally:
+            _cleanup_port_forward(f"pod/{pod}", pf)
+
+        # The drain must END in an exit. Read off `restartCount`: the process
+        # exits when the drain elapses and kubelet restarts the container in
+        # place, same pod. Watching this rather than the listener closing is
+        # deliberate — the restart is fast enough that a port-forward probe can
+        # miss the closed window entirely and hang, whereas `restartCount` is
+        # monotonic and cannot be missed. The poll's own timeout is deliberately
+        # looser than the ceiling below, so a lengthened drain fails on the
+        # assertion (which explains it) rather than on a bare TimeoutError.
+        _poll_until(
+            lambda: _container_restart_count(pod, "router") > restarts_before,
+            "router container restarts once the drain elapses",
+            timeout=CONFIGURED_DRAIN_SECS + RESTART_OBSERVATION_SLACK_SECS + 30,
+            interval=0.5,
+        )
+        # The drain's CEILING. Its mirror image — the pause not being cut short
+        # — is the still-up assertion inside the window above; together they
+        # bound the pause from both sides, which neither does alone.
+        held_open_for = time.monotonic() - sigterm_at
+        assert held_open_for < CONFIGURED_DRAIN_SECS + RESTART_OBSERVATION_SLACK_SECS, (
+            f"the router was still up {held_open_for:.1f}s after SIGTERM, past the "
+            f"configured {CONFIGURED_DRAIN_SECS}s drain by more than kubelet's restart "
+            f"latency can explain — check the seconds-to-Duration conversion in "
+            f"ServerConfig::shutdown_drain()"
+        )
+        _wait_for_deployment_ready("sgl-router")

experimental/sgl-router/tests/proxy/graceful_shutdown.rs

@@ -2,10 +2,11 @@
 // SPDX-License-Identifier: Apache-2.0
 
 //! Pins the contract that `axum::serve(...).with_graceful_shutdown(...)` —
-//! exactly as wired in `src/main.rs` — drains every in-flight streaming
+//! the same combinator `src/main.rs` uses — drains every in-flight streaming
 //! request through the **real** `build_router(ctx)` stack before the
 //! server future resolves. A k8s SIGTERM must not truncate streaming
-//! completions.
+//! completions. (`main.rs` additionally runs the readiness drain first; the
+//! later tests cover that.)
 //!
 //! Why route the test through the real router (chat handler + proxy +
 //! SSE pump) rather than a synthetic `Router::new().route(...)`: a
@@ -13,6 +14,14 @@
 //! `bytes_stream_to_body` completion hook, in `chat::chat_completions`'
 //! guards, or in the SSE pump's `tx.send().await` race — all of which
 //! would be silently skipped by a synthetic-handler test.
+//!
+//! The later tests pin the readiness drain that runs *before* that axum
+//! drain: `server::shutdown::drain_for_termination` flips `/readyz` to 503
+//! and holds the listener open for `--shutdown-drain-secs` so the endpoint
+//! removal reaches kube-proxy first. They substitute a channel for the real
+//! `Signal`, so `main.rs`'s `shutdown_signal` is not exercised here; the k8s
+//! integration suite (`tests/e2e/k8s_integration/test_shutdown_drain.py`)
+//! signals the shipped binary and covers that wiring.
 
 use futures::future::join_all;
 use sgl_router::config::{
@@ -235,6 +244,293 @@ async fn shutdown_with_no_inflight_returns_promptly() {
     );
 }
 
+/// The readiness-drain contract: on SIGTERM the drain flips `/readyz` to 503
+/// *while the server keeps accepting* (`/healthz` stays 200, a brand-new
+/// connection is still served), so the endpoint removal reaches kube-proxy
+/// before the listener closes. Mirrors `src/main.rs`'s SIGTERM arm by driving
+/// the shutdown future as "await the signal, then `drain_for_termination`"
+/// against the real `build_router(ctx)` stack.
+///
+/// The drain window is ended by the `expedite` channel rather than by wall
+/// clock, so the mid-drain assertions cannot lose a race with a sleeping
+/// timer on a loaded runner — and the expedite path itself gets covered.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn readyz_flips_to_503_during_drain_while_still_serving() {
+    let worker = crate::common::mock_worker::MockWorker::start_slow_stream(
+        SLOW_CHUNKS.to_vec(),
+        Duration::from_millis(20),
+    )
+    .await;
+    let ctx = build_ctx_with_worker(&worker.url);
+    assert!(ctx.is_ready(), "ctx starts ready");
+
+    let app = build_router(ctx.clone());
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+
+    // `sigterm_tx` stands in for SIGTERM delivery; `expedite_tx` stands in for
+    // the further termination signal that cuts the pause short. The drain is
+    // an hour so only `expedite_tx` can end it.
+    let ctx_for_shutdown = ctx.clone();
+    let (sigterm_tx, sigterm_rx) = oneshot::channel::<()>();
+    let (expedite_tx, expedite_rx) = oneshot::channel::<()>();
+    let server = tokio::spawn(async move {
+        axum::serve(listener, app)
+            .with_graceful_shutdown(async move {
+                let _ = sigterm_rx.await;
+                sgl_router::server::shutdown::drain_for_termination(
+                    &ctx_for_shutdown,
+                    Duration::from_secs(3600),
+                    async {
+                        let _ = expedite_rx.await;
+                    },
+                )
+                .await;
+            })
+            .await
+            .unwrap();
+    });
+
+    // Every probe opens its own connection: a pooled client would ride the
+    // pre-SIGTERM connection and keep passing even if the listener had already
+    // closed, which is exactly the regression this test exists to catch.
+    let client = reqwest::Client::builder()
+        .pool_max_idle_per_host(0)
+        .build()
+        .unwrap();
+    let readyz = format!("http://{addr}/readyz");
+    let healthz = format!("http://{addr}/healthz");
+
+    // Before SIGTERM: ready + worker registered ⇒ /readyz 200.
+    let pre = client.get(&readyz).send().await.unwrap();
+    assert_eq!(
+        pre.status(),
+        reqwest::StatusCode::OK,
+        "ready before SIGTERM"
+    );
+
+    sigterm_tx.send(()).unwrap();
+    // The drain flips readiness before its first await, but the flip and this
+    // observation are on different tasks — wait for it rather than sleeping.
+    tokio::time::timeout(Duration::from_secs(5), async {
+        while ctx.is_ready() {
+            tokio::time::sleep(Duration::from_millis(5)).await;
+        }
+    })
+    .await
+    .expect("the drain must flip readiness off promptly after SIGTERM");
+
+    let mid_ready = client.get(&readyz).send().await.unwrap();
+    assert_eq!(
+        mid_ready.status(),
+        reqwest::StatusCode::SERVICE_UNAVAILABLE,
+        "/readyz must flip to 503 during the drain so probes and load balancers see this pod as not-ready before the listener closes",
+    );
+
+    // State the accept explicitly rather than inferring it from a 200: this is
+    // the half of the contract that a pooled client would silently satisfy.
+    tokio::net::TcpStream::connect(addr)
+        .await
+        .expect("the listener must still accept new connections during the drain");
+    let mid_health = client.get(&healthz).send().await.unwrap();
+    assert_eq!(
+        mid_health.status(),
+        reqwest::StatusCode::OK,
+        "the server must still be serving during the drain window",
+    );
+
+    // A *real proxied* request (not just the local health handlers) must still
+    // be accepted and served during the drain window — this is the request k8s
+    // may still route before the endpoint removal reaches kube-proxy.
+    let chat = format!("http://{addr}/v1/chat/completions");
+    let body = serde_json::json!({
+        "model": "tiny",
+        "messages": [{"role": "user", "content": "hi"}],
+    });
+    let mid_chat = client.post(&chat).json(&body).send().await.unwrap();
+    assert_eq!(
+        mid_chat.status(),
+        reqwest::StatusCode::OK,
+        "a proxied chat request must still succeed during the drain window",
+    );
+
+    // The request the drain actually exists for: it ARRIVES during the pause
+    // (kube-proxy has not observed the removal yet) and is still streaming when
+    // the pause ends. It must survive the handover into axum's in-flight drain,
+    // not just the window it started in.
+    //
+    // Await the response headers here rather than inside the spawned task: that
+    // is the point at which the request is provably in flight, so cutting the
+    // pause short below cannot race the client's connect on a loaded runner.
+    let stream_client = reqwest::Client::builder()
+        .timeout(Duration::from_secs(10))
+        .build()
+        .unwrap();
+    let late_request = serde_json::json!({
+        "model": "tiny",
+        "messages": [{"role": "user", "content": "hi"}],
+        "stream": true,
+    });
+    let late_resp = stream_client
+        .post(&chat)
+        .json(&late_request)
+        .send()
+        .await
+        .unwrap();
+    assert!(
+        late_resp.status().is_success(),
+        "a stream started during the drain must be accepted: {}",
+        late_resp.status(),
+    );
+    let late = tokio::spawn(async move { late_resp.bytes().await.unwrap() });
+
+    // Cut the pause short while that stream is still mid-flight; the server
+    // resolves without waiting out the hour.
+    expedite_tx.send(()).unwrap();
+
+    let late_body = late.await.expect("late client task joined");
+    assert!(
+        String::from_utf8_lossy(&late_body).contains("data: [DONE]"),
+        "a request that arrived during the drain must still complete after the pause ends",
+    );
+    tokio::time::timeout(Duration::from_secs(5), server)
+        .await
+        .expect("an expedite signal must end the drain instead of sleeping an hour")
+        .expect("server task joined cleanly");
+}
+
+/// After the drain elapses and the server future resolves, axum must have
+/// stopped accepting: a *new* connection is refused. This is the other half of
+/// the contract — the drain has to actually END in a closed listener, or the
+/// pause merely postpones shutdown without ever handing traffic off. (What
+/// closes the rolling-update race is the pause itself, covered by
+/// `readyz_flips_to_503_during_drain_while_still_serving`.) Asserted on a raw
+/// TCP connect so the failure has to be `ConnectionRefused`; a `reqwest` error
+/// would also cover a timeout, which is a different (and on a loaded runner,
+/// plausible) outcome.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn new_connections_refused_after_drain_completes() {
+    let worker = crate::common::mock_worker::MockWorker::start(vec![]).await;
+    let ctx = build_ctx_with_worker(&worker.url);
+    let app = build_router(ctx.clone());
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+
+    // Short drain so the test is fast; the point is the post-resolve state.
+    let drain = Duration::from_millis(100);
+    let ctx_for_shutdown = ctx.clone();
+    let (sigterm_tx, sigterm_rx) = oneshot::channel::<()>();
+    let server = tokio::spawn(async move {
+        axum::serve(listener, app)
+            .with_graceful_shutdown(async move {
+                let _ = sigterm_rx.await;
+                sgl_router::server::shutdown::drain_for_termination(
+                    &ctx_for_shutdown,
+                    drain,
+                    std::future::pending::<()>(),
+                )
+                .await;
+            })
+            .await
+            .unwrap();
+    });
+
+    // Server accepts before shutdown.
+    tokio::net::TcpStream::connect(addr)
+        .await
+        .expect("listener accepts before SIGTERM");
+
+    // Fire SIGTERM and wait for the drain + server future to fully resolve.
+    sigterm_tx.send(()).unwrap();
+    tokio::time::timeout(Duration::from_secs(5), server)
+        .await
+        .expect("server resolves after the drain elapses")
+        .expect("server task joined cleanly");
+
+    // A fresh connection must now be refused — the listener is closed.
+    let err = tokio::net::TcpStream::connect(addr)
+        .await
+        .expect_err("a new connection must be refused after the drain completes");
+    assert_eq!(
+        err.kind(),
+        std::io::ErrorKind::ConnectionRefused,
+        "expected the closed listener to refuse, got {err:?}",
+    );
+}
+
+/// End-to-end composition: SIGTERM → `drain_for_termination` (flip 503, pause)
+/// → axum drains the already-attached streaming request to `[DONE]`.
+/// `shutdown_drains_100_inflight_streaming_chat_completions` drives a bare
+/// oneshot shutdown future; this one composes the readiness drain with the axum
+/// drain, so a regression that truncates in-flight streams once the drain
+/// begins is caught. It does NOT assert the flip/pause ordering —
+/// `readyz_flips_to_503_during_drain_while_still_serving` covers that.
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn inflight_stream_completes_through_drain_for_termination() {
+    let worker = crate::common::mock_worker::MockWorker::start_slow_stream(
+        SLOW_CHUNKS.to_vec(),
+        Duration::from_millis(60),
+    )
+    .await;
+    let ctx = build_ctx_with_worker(&worker.url);
+    let app = build_router(ctx.clone());
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    let url = format!("http://{addr}/v1/chat/completions");
+
+    let drain = Duration::from_millis(50);
+    let ctx_for_shutdown = ctx.clone();
+    let (sigterm_tx, sigterm_rx) = oneshot::channel::<()>();
+    let server = tokio::spawn(async move {
+        axum::serve(listener, app)
+            .with_graceful_shutdown(async move {
+                let _ = sigterm_rx.await;
+                sgl_router::server::shutdown::drain_for_termination(
+                    &ctx_for_shutdown,
+                    drain,
+                    std::future::pending::<()>(),
+                )
+                .await;
+            })
+            .await
+            .unwrap();
+    });
+
+    // Start one slow stream and hand back the response only once its headers
+    // have arrived — that is the point at which the request is provably
+    // in-flight, so SIGTERM below cannot race the client's connect.
+    let client = reqwest::Client::builder()
+        .timeout(Duration::from_secs(10))
+        .build()
+        .unwrap();
+    let body = serde_json::json!({
+        "model": "tiny",
+        "messages": [{"role": "user", "content": "hi"}],
+        "stream": true,
+    });
+    let resp = client.post(&url).json(&body).send().await.unwrap();
+    assert!(
+        resp.status().is_success(),
+        "stream started: {}",
+        resp.status()
+    );
+    let inflight = tokio::spawn(async move { resp.bytes().await.unwrap() });
+
+    // Fire SIGTERM mid-stream: the drain must NOT truncate the in-flight stream.
+    sigterm_tx.send(()).unwrap();
+
+    let received = inflight.await.expect("client task joined");
+    let body_str = String::from_utf8_lossy(&received);
+    assert!(
+        body_str.contains("data: [DONE]"),
+        "the in-flight stream must terminate with `data: [DONE]` through the drain path, got: {body_str}",
+    );
+    tokio::time::timeout(Duration::from_secs(5), server)
+        .await
+        .expect("server resolves after in-flight stream drains")
+        .expect("server task joined cleanly");
+}
+
 /// Poll until `inflight_http` settles on `want`, so the assertions below do not
 /// race the guard drop that happens on the server task after the client has
 /// already seen the last byte.