[sgl-router] Bound streaming lifetimes and release guards on idle disconnect (#40391)

Co-authored-by: Shangming Cai <csmthu@gmail.com>

experimental/sgl-router/src/config/cli.rs

@@ -114,9 +114,13 @@ pub struct ServerArgs {
     /// Per-request upstream timeout in seconds.
     #[arg(long, default_value_t = default_proxy_request_timeout_secs())]
     pub request_timeout_secs: u64,
+    /// Maximum silence between upstream stream chunks, in seconds.
+    #[arg(long, default_value_t = ProxyConfig::default().stream_idle_timeout_secs)]
+    pub stream_idle_timeout_secs: u64,
 
-    /// Max lifetime of an in-flight request entry before the janitor
-    /// reaps it (returns 504 `stale_request_expired`).
+    /// Maximum in-flight request lifetime in seconds, including streaming responses.
+    /// Expiry returns 504 `stale_request_expired` before response headers are sent;
+    /// after streaming starts, it aborts the body without changing the HTTP status.
     #[arg(long, default_value_t = default_stale_request_timeout_secs())]
     pub stale_request_timeout_secs: u64,
 
@@ -389,6 +393,7 @@ impl Cli {
             discovery,
             proxy: ProxyConfig {
                 request_timeout_secs: self.server.request_timeout_secs,
+                stream_idle_timeout_secs: self.server.stream_idle_timeout_secs,
             },
             router_inflight_load: InflightLoadConfig {
                 stale_request_timeout_secs: self.server.stale_request_timeout_secs,

experimental/sgl-router/src/config/mod.rs

@@ -48,6 +48,10 @@ impl Config {
             validate_bucket_config(bucket_config)?;
         }
         self.model.sampling_overrides.validate()?;
+        ensure!(
+            self.proxy.stream_idle_timeout_secs > 0,
+            "stream_idle_timeout_secs must be greater than zero"
+        );
         ensure!(
             self.server.shutdown_drain_secs <= MAX_SHUTDOWN_DRAIN_SECS,
             "shutdown_drain_secs must be at most {MAX_SHUTDOWN_DRAIN_SECS} (got {}); \

experimental/sgl-router/src/config/types.rs

@@ -19,6 +19,8 @@ pub struct Config {
 pub struct ProxyConfig {
     /// Timeout for upstream response headers and body. Counts as a circuit-breaker failure.
     pub request_timeout_secs: u64,
+    /// Maximum silence between streamed upstream chunks before the stream fails.
+    pub stream_idle_timeout_secs: u64,
 }
 
 pub fn default_proxy_request_timeout_secs() -> u64 {
@@ -29,6 +31,7 @@ impl Default for ProxyConfig {
     fn default() -> Self {
         Self {
             request_timeout_secs: default_proxy_request_timeout_secs(),
+            stream_idle_timeout_secs: 180,
         }
     }
 }

experimental/sgl-router/src/main.rs

@@ -276,7 +276,8 @@ fn build_app_context(
     let block_size_oracle = engine_state.block_size_oracle();
     let proxy = Arc::new(
         Proxy::new(Duration::from_secs(config.proxy.request_timeout_secs))
-            .context("build proxy client")?,
+            .context("build proxy client")?
+            .with_stream_idle_timeout(Duration::from_secs(config.proxy.stream_idle_timeout_secs)),
     );
 
     let mut app_context = AppContext::with_router_inflight_load(

experimental/sgl-router/src/proxy/mod.rs

@@ -16,6 +16,7 @@ use bytes::Bytes;
 use reqwest::{Client, Url};
 use std::sync::Arc;
 use std::time::Duration;
+use tokio_util::sync::CancellationToken;
 
 /// Parse a worker URL emitted by discovery.  On failure, trip the worker's
 /// circuit breaker so the malformed worker drops out of subsequent
@@ -38,14 +39,14 @@ enum BreakerOutcome {
     /// The worker was responsive — a 2xx, or a 4xx it answered cleanly (a
     /// client's bad request says nothing about worker health). The non-streaming
     /// arm records success immediately; the streaming arm defers to the pump's
-    /// completion hook, which records success or failure by
-    /// [`sse::StreamEnd::transport_ok`], since a 2xx head can still be followed
+    /// completion hook, which classifies [`sse::StreamEnd::reason`],
+    /// since a 2xx head can still be followed
     /// by a body that never completes.
     Success,
     /// A real fault (5xx other than backpressure) → `record_failure`: count
     /// toward opening.
     Failure,
-    /// Backpressure (the worker is responsive but at capacity) →
+    /// Backpressure or router-side stream expiry →
     /// `record_backpressure`: never opens the breaker and, while Closed, leaves
     /// an in-progress failure streak intact — but still resolves a half-open
     /// probe so a recovered-but-busy worker isn't wedged shut.
@@ -81,6 +82,20 @@ fn breaker_outcome(status: reqwest::StatusCode) -> BreakerOutcome {
     }
 }
 
+/// Router-side expiry says nothing about worker health. Preserve the existing
+/// treatment of completed streams and client disconnects; upstream faults,
+/// idle timeouts, and pump panics remain failures.
+fn stream_breaker_outcome(end: sse::StreamEnd) -> BreakerOutcome {
+    use sse::StreamEndReason;
+    match end.reason {
+        StreamEndReason::Expired => BreakerOutcome::Neutral,
+        StreamEndReason::Completed | StreamEndReason::ClientDisconnect => BreakerOutcome::Success,
+        StreamEndReason::UpstreamError
+        | StreamEndReason::IdleTimeout
+        | StreamEndReason::PumpPanicked => BreakerOutcome::Failure,
+    }
+}
+
 #[derive(Debug)]
 pub struct Proxy {
     /// The negotiating client: HTTP/1.1 in cleartext, and ALPN `h2, http/1.1`
@@ -94,6 +109,8 @@ pub struct Proxy {
     /// Wall-clock timeout applied to non-streaming upstream requests. Streaming
     /// requests deliberately do not use this (long generations are valid).
     pub request_timeout: Duration,
+    /// Maximum silence between streamed upstream chunks; `None` waits forever.
+    pub stream_idle_timeout: Option<Duration>,
 }
 
 /// Build a forwarding client for `protocol`, sharing pool/connect tuning
@@ -127,9 +144,15 @@ impl Proxy {
             default_client: build_client(WireProtocol::Http1)?,
             h2c_client: build_client(WireProtocol::H2c)?,
             request_timeout,
+            stream_idle_timeout: None,
         })
     }
 
+    pub fn with_stream_idle_timeout(mut self, timeout: Duration) -> Self {
+        self.stream_idle_timeout = Some(timeout);
+        self
+    }
+
     /// The forwarding client for `protocol`, taken from the selected worker's
     /// [`crate::workers::Worker::protocol`].
     fn client_for(&self, protocol: WireProtocol) -> &Client {
@@ -282,6 +305,7 @@ impl Proxy {
         stream_guards: Option<Box<dyn Send + 'static>>,
         on_first_byte: Option<Box<dyn FnOnce() + Send + 'static>>,
         on_stream_end: Option<Box<dyn FnOnce(sse::StreamEnd) + Send + 'static>>,
+        expiration: Option<CancellationToken>,
     ) -> Result<Response<Body>, ApiError> {
         if !breaker.allow() {
             return Err(ApiError::BreakerOpen {
@@ -345,10 +369,10 @@ impl Proxy {
                 BreakerOutcome::Success => {
                     let breaker_for_hook = Arc::clone(breaker);
                     Some(Box::new(move |end| {
-                        if end.transport_ok {
-                            breaker_for_hook.record_success();
-                        } else {
-                            breaker_for_hook.record_failure();
+                        match stream_breaker_outcome(end) {
+                            BreakerOutcome::Success => breaker_for_hook.record_success(),
+                            BreakerOutcome::Failure => breaker_for_hook.record_failure(),
+                            BreakerOutcome::Neutral => breaker_for_hook.record_backpressure(),
                         }
                         if let Some(hook) = caller_end_hook {
                             hook(end);
@@ -368,6 +392,10 @@ impl Proxy {
             stream_guards,
             on_complete,
             first_byte_hook,
+            sse::StreamLimits {
+                idle_timeout: self.stream_idle_timeout,
+                expiration,
+            },
         );
         let mut out = Response::new(body);
         *out.status_mut() = status;
@@ -467,6 +495,145 @@ mod tests {
         (format!("http://127.0.0.1:{port}"), tx)
     }
 
+    async fn spawn_pending_stream_worker() -> (String, oneshot::Sender<()>) {
+        use futures::StreamExt;
+
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+        let app = Router::new().route(
+            "/v1/chat/completions",
+            post(|| async {
+                Body::from_stream(
+                    futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(
+                        b"data: chunk\n\n",
+                    ))])
+                    .chain(futures::stream::pending()),
+                )
+            }),
+        );
+        let (tx, rx) = oneshot::channel();
+        tokio::spawn(async move {
+            axum::serve(listener, app)
+                .with_graceful_shutdown(async move {
+                    let _ = rx.await;
+                })
+                .await
+                .unwrap();
+        });
+        (format!("http://{address}"), tx)
+    }
+
+    async fn pending_stream_body(
+        proxy: &Proxy,
+        url: &str,
+        breaker: &Arc<CircuitBreaker>,
+        expiration: Option<CancellationToken>,
+    ) -> Body {
+        use http_body_util::BodyExt;
+
+        let response = proxy
+            .forward_streaming_to(
+                url,
+                WireProtocol::Http1,
+                breaker,
+                "/v1/chat/completions",
+                &HeaderMap::new(),
+                Bytes::from_static(b"{}"),
+                None,
+                None,
+                None,
+                expiration,
+            )
+            .await
+            .unwrap();
+        assert_eq!(response.status(), StatusCode::OK);
+        let mut body = response.into_body();
+        tokio::time::timeout(Duration::from_secs(2), body.frame())
+            .await
+            .unwrap()
+            .unwrap()
+            .unwrap();
+        body
+    }
+
+    #[tokio::test]
+    async fn stream_expiry_preserves_breaker_failure_streak() {
+        use http_body_util::BodyExt;
+
+        let (url, _shutdown) = spawn_pending_stream_worker().await;
+        let proxy = Proxy::new(Duration::from_secs(5)).unwrap();
+        let breaker = Arc::new(CircuitBreaker::new());
+        breaker.record_failure();
+        breaker.record_failure();
+        for _ in 0..6 {
+            let expiration = CancellationToken::new();
+            let body = pending_stream_body(&proxy, &url, &breaker, Some(expiration.clone())).await;
+            expiration.cancel();
+            let error = tokio::time::timeout(Duration::from_secs(2), body.collect())
+                .await
+                .unwrap()
+                .unwrap_err();
+            assert!(error.to_string().contains("stale_request_timeout"));
+            assert_eq!(
+                breaker.snapshot().state_code,
+                0,
+                "expiry must not add a failure"
+            );
+        }
+        breaker.record_failure();
+        assert_eq!(
+            breaker.snapshot().state_code,
+            1,
+            "expiry must not reset prior failures"
+        );
+    }
+
+    #[tokio::test]
+    async fn stream_expiry_resolves_half_open_probe() {
+        use http_body_util::BodyExt;
+
+        let (url, _shutdown) = spawn_pending_stream_worker().await;
+        let proxy = Proxy::new(Duration::from_secs(5)).unwrap();
+        let breaker = Arc::new(CircuitBreaker::with_config(CircuitBreakerConfig {
+            threshold: NonZeroU32::new(1).unwrap(),
+            cool_down: Duration::ZERO,
+        }));
+        breaker.record_failure();
+        let expiration = CancellationToken::new();
+        let body = pending_stream_body(&proxy, &url, &breaker, Some(expiration.clone())).await;
+        assert_eq!(breaker.snapshot().state_code, 2);
+        assert!(!breaker.would_allow());
+        expiration.cancel();
+        let error = tokio::time::timeout(Duration::from_secs(2), body.collect())
+            .await
+            .unwrap()
+            .unwrap_err();
+        assert!(error.to_string().contains("stale_request_timeout"));
+        assert_eq!(breaker.snapshot().state_code, 0);
+        assert!(breaker.would_allow());
+    }
+
+    #[tokio::test]
+    async fn stream_idle_timeout_still_trips_breaker() {
+        use http_body_util::BodyExt;
+
+        let (url, _shutdown) = spawn_pending_stream_worker().await;
+        let mut proxy = Proxy::new(Duration::from_secs(5)).unwrap();
+        proxy.stream_idle_timeout = Some(Duration::from_millis(20));
+        let breaker = Arc::new(CircuitBreaker::with_config(CircuitBreakerConfig {
+            threshold: NonZeroU32::new(1).unwrap(),
+            cool_down: Duration::from_secs(30),
+        }));
+        let body = pending_stream_body(&proxy, &url, &breaker, None).await;
+        let error = tokio::time::timeout(Duration::from_secs(2), body.collect())
+            .await
+            .unwrap()
+            .unwrap_err();
+        assert!(error.to_string().contains("idle timeout"));
+        assert_eq!(breaker.snapshot().state_code, 1);
+        assert!(!breaker.would_allow());
+    }
+
     /// A saturated engine's own queue-full 503s must not trip the router's
     /// circuit breaker. Dispatch far past any plausible failure threshold and
     /// assert the breaker stays Closed and admitting.
@@ -610,6 +777,7 @@ mod tests {
                     None,
                     None,
                     None,
+                    None,
                 )
                 .await
                 .expect("streaming dispatch should reach the worker");

experimental/sgl-router/src/proxy/sse.rs

@@ -4,22 +4,32 @@
 //! SSE passthrough — bridges a reqwest `bytes_stream()` into an axum Body.
 
 use std::panic::AssertUnwindSafe;
-use std::sync::Arc;
+use std::time::Duration;
 
 use axum::body::Body;
 use bytes::Bytes;
 use futures::{FutureExt, StreamExt};
 use tokio_stream::wrappers::ReceiverStream;
+use tokio_util::sync::CancellationToken;
+
+/// Why the SSE pump stopped, independently of any SSE error event it observed.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum StreamEndReason {
+    Completed,
+    UpstreamError,
+    IdleTimeout,
+    /// The router's stale-request deadline expired, regardless of worker health.
+    Expired,
+    ClientDisconnect,
+    PumpPanicked,
+}
 
 /// How the SSE pump ended, reported to the `on_complete` hook.
 #[derive(Debug, Clone, Copy)]
 pub struct StreamEnd {
-    /// No upstream stream error and no pump panic.
-    pub transport_ok: bool,
+    pub reason: StreamEndReason,
     /// An SSE error event (`data: {"error"...}`) rode the stream.
     pub saw_error_event: bool,
-    /// The client dropped the response body before upstream finished.
-    pub client_disconnect: bool,
 }
 
 /// A `data:` line whose payload's first JSON key is `error` — tolerant of
@@ -59,140 +69,126 @@ impl ErrorEventScanner {
     }
 }
 
+/// Bounds on a streaming response beyond what the upstream stream itself provides.
+#[derive(Debug, Clone, Default)]
+pub struct StreamLimits {
+    /// Maximum silence between upstream chunks. `None` waits indefinitely.
+    pub idle_timeout: Option<Duration>,
+    /// Fires when the stale-request janitor expires the request.
+    pub expiration: Option<CancellationToken>,
+}
+
 /// Bridge a byte stream into an axum Body that streams chunks unchanged.
 ///
-/// Spawns one tokio task per stream so the handler can return immediately.
-/// Uses a **bounded** 64-slot channel so `tx.send().await` naturally
-/// backpressures the upstream read when the client (axum Body consumer) falls
-/// behind — an unbounded channel would buffer hundreds of MB for a slow client
-/// receiving a long completion.
+/// One tokio task pumps upstream chunks through a bounded 64-slot channel so a
+/// slow client backpressures the upstream read. The pump stops as soon as the
+/// client disconnects, even while upstream is silent, when `limits.idle_timeout`
+/// elapses between chunks, or when `limits.expiration` fires.
 ///
-/// # Backpressure note
-/// The channel bound of 64 absorbs short bursts while still limiting
-/// worst-case outstanding bytes to 64 × chunk_size (typically a few MB).
+/// The terminal result travels on a separate channel and is chained after the
+/// data, so a full queue cannot block cleanup or turn a failed stream into a
+/// clean EOF. A pump panic is reported the same way.
 ///
-/// # Client disconnect
-/// When the axum Body is dropped the receiver is closed; `tx.send()` then
-/// returns `Err`, which breaks the loop — no upstream bytes are read after the
-/// client disconnects.
-///
-/// # Panic safety
-/// The pump future is wrapped in `AssertUnwindSafe(..).catch_unwind()`. If the
-/// upstream stream panics, we surface a loud `io::Error` to the client; without
-/// this, the body would EOF cleanly and clients couldn't distinguish that from
-/// success — the worst failure class (truncated output that looks complete).
-///
-/// # Stream guards
-/// When `stream_guards` is `Some`, the value is **moved into the spawned task**
-/// and held for the entire body lifetime.  It is dropped only when the SSE
-/// pump finishes (stream exhausted, client disconnects, or upstream errors).
-/// The opaque `Box<dyn Send + 'static>` accepts any drop-only payload — most
-/// commonly a tuple of [`crate::workers::LoadGuard`] and
-/// [`crate::state::load_monitor::router_inflight_load::RouterInflightLoadGuard`]. The proxy does not
-/// inspect the value; it relies entirely on `Drop` semantics, so callers can
-/// pack arbitrary cleanup state in. Pass `None` for callers that manage the
-/// guard externally (e.g. non-streaming paths where the handler itself is the
-/// guard scope).
-///
-/// # Completion hook
-/// When `on_complete` is `Some`, it runs exactly once when the pump task
-/// finishes with the transport, SSE error-event, and client-disconnect state.
-///
-/// # First-byte hook
-/// When `on_first_byte` is `Some`, the closure runs exactly once, the moment
-/// the first `Ok` chunk is read from the upstream stream — i.e. time to first
-/// token. It does NOT fire if the stream ends or errors before any `Ok` chunk
-/// arrives. `forward_streaming_to` passes a closure that records
-/// `sgl_router_ttft_seconds` for successful streaming responses.
+/// `guards` is held until the pump finishes; `on_first_byte` runs on the first
+/// `Ok` chunk; `on_complete` runs exactly once with the final [`StreamEnd`].
 pub fn bytes_stream_to_body<S, E>(
-    stream: S,
-    stream_guards: Option<Box<dyn Send + 'static>>,
+    mut stream: S,
+    guards: Option<Box<dyn Send + 'static>>,
     on_complete: Option<Box<dyn FnOnce(StreamEnd) + Send + 'static>>,
-    on_first_byte: Option<Box<dyn FnOnce() + Send + 'static>>,
+    mut on_first_byte: Option<Box<dyn FnOnce() + Send + 'static>>,
+    limits: StreamLimits,
 ) -> Body
 where
     S: futures::Stream<Item = Result<Bytes, E>> + Send + Unpin + 'static,
     E: std::fmt::Display + Send + Sync + 'static,
 {
-    let (tx, rx) = tokio::sync::mpsc::channel(64);
+    let (tx, rx) = tokio::sync::mpsc::channel::<Bytes>(64);
+    let (terminal_tx, terminal_rx) = tokio::sync::oneshot::channel();
     tokio::spawn(async move {
-        let tx_for_panic = tx.clone();
-        let outcome = Arc::new(parking_lot::Mutex::new(StreamEnd {
-            transport_ok: true,
+        let mut end = StreamEnd {
+            reason: StreamEndReason::Completed,
             saw_error_event: false,
-            client_disconnect: false,
-        }));
-        let outcome_setter = Arc::clone(&outcome);
-        // `None` once an error event is found — the scan is done for good.
-        let mut scanner = Some(ErrorEventScanner::default());
-        let pump = AssertUnwindSafe(async move {
-            // Hold the guards for the task's lifetime — dropped when this
-            // block exits (stream done or client disconnect).  Leading
-            // underscore suppresses the "unused variable" lint while
-            // keeping intent explicit.
-            let _hold = stream_guards;
-            let mut on_first_byte = on_first_byte;
-            let mut s = stream;
-            while let Some(chunk) = s.next().await {
-                let item: Result<Bytes, std::io::Error> = chunk.map_err(|e| {
-                    let msg = e.to_string();
-                    tracing::warn!(error = %msg, "upstream SSE stream errored mid-flight");
-                    std::io::Error::other(msg)
-                });
-                let is_err_chunk = item.is_err();
-                match &item {
-                    Ok(bytes) => {
-                        // TTFT hook: at most once (`take()`); an error-first
-                        // stream never produced a token, so it stays unfired.
+        };
+        let mut scanner = ErrorEventScanner::default();
+        let idle = limits.idle_timeout.unwrap_or(Duration::MAX);
+        let expired = async {
+            match limits.expiration {
+                Some(token) => token.cancelled().await,
+                None => std::future::pending().await,
+            }
+        };
+        // Disconnect and expiration race the whole forwarding loop, so they
+        // fire while `send` waits on a full queue as well as while upstream is silent.
+        let pump = async {
+            tokio::select! {
+                biased;
+                _ = tx.closed() => {
+                    end.reason = StreamEndReason::ClientDisconnect;
+                    Ok(())
+                }
+                _ = expired => {
+                    end.reason = StreamEndReason::Expired;
+                    Err(std::io::Error::other("SSE stream exceeded stale_request_timeout"))
+                }
+                result = async {
+                    loop {
+                        let bytes = match tokio::time::timeout(idle, stream.next()).await {
+                            Ok(None) => return (StreamEndReason::Completed, Ok(())),
+                            Ok(Some(Ok(bytes))) => bytes,
+                            Ok(Some(Err(e))) => return (
+                                StreamEndReason::UpstreamError,
+                                Err(std::io::Error::other(e.to_string())),
+                            ),
+                            Err(_) => return (
+                                StreamEndReason::IdleTimeout,
+                                Err(std::io::Error::other("SSE upstream idle timeout")),
+                            ),
+                        };
                         if let Some(hook) = on_first_byte.take() {
                             hook();
                         }
-                        if scanner.as_mut().is_some_and(|scanner| scanner.feed(bytes)) {
-                            outcome_setter.lock().saw_error_event = true;
-                            scanner = None;
+                        if !end.saw_error_event {
+                            end.saw_error_event = scanner.feed(&bytes);
+                        }
+                        if tx.send(bytes).await.is_err() {
+                            // Receiver gone; the `tx.closed()` arm reports the disconnect.
+                            std::future::pending::<()>().await;
                         }
                     }
-                    Err(_) => outcome_setter.lock().transport_ok = false,
-                }
-                if tx.send(item).await.is_err() {
-                    // Receiver dropped. If we were about to ship an upstream
-                    // error there's nothing left to report; otherwise this is
-                    // a clean client-side disconnect — log at debug since it's
-                    // not a router-side fault.
-                    if !is_err_chunk {
-                        tracing::debug!("SSE client disconnected mid-stream");
-                        outcome_setter.lock().client_disconnect = true;
-                    }
-                    break;
-                }
-                if is_err_chunk {
-                    // Surfaced upstream error to client; stop reading.
-                    break;
+                } => {
+                    end.reason = result.0;
+                    result.1
                 }
             }
-        });
-        let pump_result = pump.catch_unwind().await;
-        let panicked = pump_result.is_err();
-        if let Err(panic_payload) = pump_result {
-            let msg = panic_payload
-                .downcast_ref::<&'static str>()
-                .map(|s| (*s).to_string())
-                .or_else(|| panic_payload.downcast_ref::<String>().cloned())
-                .unwrap_or_else(|| "<non-string panic payload>".to_string());
-            tracing::error!(error = %msg, "SSE pump task panicked");
-            let _ = tx_for_panic
-                .send(Err(std::io::Error::other(format!(
-                    "SSE pump panicked: {msg}"
-                ))))
-                .await;
+        };
+        let result = match AssertUnwindSafe(pump).catch_unwind().await {
+            Ok(result) => result,
+            Err(payload) => {
+                end.reason = StreamEndReason::PumpPanicked;
+                let message = payload
+                    .downcast_ref::<&str>()
+                    .copied()
+                    .or_else(|| payload.downcast_ref::<String>().map(String::as_str))
+                    .unwrap_or("<non-string panic payload>");
+                Err(std::io::Error::other(format!(
+                    "SSE pump panicked: {message}"
+                )))
             }
+        };
         if let Some(hook) = on_complete {
-            let mut end = *outcome.lock();
-            end.transport_ok &= !panicked;
             hook(end);
         }
+        drop(guards);
+        let _ = terminal_tx.send(result);
     });
-    Body::from_stream(ReceiverStream::new(rx))
+    let terminal = futures::stream::once(terminal_rx).filter_map(|result| {
+        futures::future::ready(match result {
+            Ok(Ok(())) => None,
+            Ok(Err(error)) => Some(Err(error)),
+            Err(_) => Some(Err(std::io::Error::other("SSE pump cancelled"))),
+        })
+    });
+    Body::from_stream(ReceiverStream::new(rx).map(Ok).chain(terminal))
 }
 
 #[cfg(test)]
@@ -202,6 +198,108 @@ mod tests {
     use futures::stream;
     use http_body_util::BodyExt;
 
+    fn limited_body<S>(
+        stream: S,
+        limits: StreamLimits,
+    ) -> (Body, tokio::sync::oneshot::Receiver<StreamEnd>)
+    where
+        S: futures::Stream<Item = Result<Bytes, std::io::Error>> + Send + Unpin + 'static,
+    {
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        let body = bytes_stream_to_body(
+            stream,
+            None,
+            Some(Box::new(move |end| {
+                let _ = tx.send(end);
+            })),
+            None,
+            limits,
+        );
+        (body, rx)
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn idle_disconnect_releases_guards_without_waiting_for_upstream() {
+        struct Release(Option<tokio::sync::oneshot::Sender<()>>);
+        impl Drop for Release {
+            fn drop(&mut self) {
+                let _ = self.0.take().unwrap().send(());
+            }
+        }
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        let body = bytes_stream_to_body(
+            stream::pending::<Result<Bytes, std::io::Error>>(),
+            Some(Box::new(Release(Some(tx)))),
+            None,
+            None,
+            StreamLimits::default(),
+        );
+        tokio::task::yield_now().await;
+        drop(body);
+        tokio::time::timeout(Duration::from_millis(1), rx)
+            .await
+            .unwrap()
+            .unwrap();
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn idle_timeout_is_a_visible_upstream_failure() {
+        let (body, end) = limited_body(
+            stream::pending(),
+            StreamLimits {
+                idle_timeout: Some(Duration::from_secs(1)),
+                expiration: None,
+            },
+        );
+        assert!(body
+            .collect()
+            .await
+            .unwrap_err()
+            .to_string()
+            .contains("idle timeout"));
+        let end = end.await.unwrap();
+        assert_eq!(end.reason, StreamEndReason::IdleTimeout);
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn expiration_releases_guards_while_queue_is_full() {
+        struct Release(Option<tokio::sync::oneshot::Sender<()>>);
+        impl Drop for Release {
+            fn drop(&mut self) {
+                let _ = self.0.take().unwrap().send(());
+            }
+        }
+        let token = CancellationToken::new();
+        let (release_tx, release_rx) = tokio::sync::oneshot::channel();
+        let (end_tx, end_rx) = tokio::sync::oneshot::channel();
+        let body = bytes_stream_to_body(
+            stream::repeat_with(|| Ok::<_, std::io::Error>(Bytes::from_static(b"chunk"))),
+            Some(Box::new(Release(Some(release_tx)))),
+            Some(Box::new(move |end| {
+                let _ = end_tx.send(end);
+            })),
+            None,
+            StreamLimits {
+                idle_timeout: None,
+                expiration: Some(token.clone()),
+            },
+        );
+        tokio::task::yield_now().await;
+        token.cancel();
+        // Cleanup must finish before the client frees any channel capacity.
+        tokio::time::timeout(Duration::from_millis(1), release_rx)
+            .await
+            .expect("expiration must release guards while the queue remains full")
+            .unwrap();
+        assert_eq!(end_rx.await.unwrap().reason, StreamEndReason::Expired);
+        // The queue is full, so the failure must ride the terminal channel.
+        assert!(body
+            .collect()
+            .await
+            .unwrap_err()
+            .to_string()
+            .contains("stale_request_timeout"));
+    }
     #[tokio::test]
     async fn passes_through_a_simple_byte_stream() {
         let chunks = vec![
@@ -209,7 +307,7 @@ mod tests {
             Ok(Bytes::from_static(b"world")),
         ];
         let s = stream::iter(chunks);
-        let body = bytes_stream_to_body(s, None, None, None);
+        let body = bytes_stream_to_body(s, None, None, None, StreamLimits::default());
         let bytes = body.collect().await.unwrap().to_bytes();
         assert_eq!(&bytes[..], b"hello world");
     }
@@ -233,6 +331,7 @@ mod tests {
             Some(Box::new(move || {
                 fired_c.fetch_add(1, Ordering::SeqCst);
             })),
+            StreamLimits::default(),
         );
         let _ = body.collect().await.unwrap();
         assert_eq!(
@@ -260,6 +359,7 @@ mod tests {
             Some(Box::new(move || {
                 fired_c.fetch_add(1, Ordering::SeqCst);
             })),
+            StreamLimits::default(),
         );
         let _ = body.collect().await;
         assert_eq!(
@@ -276,7 +376,7 @@ mod tests {
             Err(std::io::Error::other("upstream blew up mid-stream")),
         ];
         let s = stream::iter(chunks);
-        let body = bytes_stream_to_body(s, None, None, None);
+        let body = bytes_stream_to_body(s, None, None, None, StreamLimits::default());
         // Collecting a body that terminates with an error must return Err.
         let result = body.collect().await;
         assert!(
@@ -338,8 +438,9 @@ mod tests {
         // that arm, the closure unwrap-or-elses would panic itself or
         // produce an empty message, which this test catches.
         let s = PanicAnyOnSecondPoll { polls: 0 };
-        let body = bytes_stream_to_body(s, None, None, None);
+        let (body, end) = limited_body(s, StreamLimits::default());
         let result = body.collect().await;
+        assert_eq!(end.await.unwrap().reason, StreamEndReason::PumpPanicked);
         assert!(
             result.is_err(),
             "expected body collect to surface non-string panic as Err, got Ok"
@@ -361,7 +462,7 @@ mod tests {
         // The pump task panics mid-stream. The client must see a loud Err,
         // NOT a silently-truncated success.
         let s = PanicOnSecondPoll { polls: 0 };
-        let body = bytes_stream_to_body(s, None, None, None);
+        let body = bytes_stream_to_body(s, None, None, None, StreamLimits::default());
         let result = body.collect().await;
         assert!(
             result.is_err(),
@@ -420,7 +521,7 @@ mod tests {
             yielded: 0,
             max: 1000, // way more than we'll let it consume
         };
-        let body = bytes_stream_to_body(stream, None, None, None);
+        let body = bytes_stream_to_body(stream, None, None, None, StreamLimits::default());
 
         // Read exactly one frame, then drop the body to simulate client disconnect.
         let mut data_stream = body.into_data_stream();
@@ -431,7 +532,7 @@ mod tests {
         // Give the pump generous time to make additional polls if its break is
         // broken. Healthy code: pump fills the 64-slot channel, then on the
         // next iteration tx.send().await detects receiver-drop and breaks.
-        tokio::time::sleep(std::time::Duration::from_millis(200)).await;
+        tokio::time::sleep(Duration::from_millis(200)).await;
         let final_polls = polls.load(Ordering::SeqCst);
         assert!(
             final_polls <= 70,
@@ -503,6 +604,7 @@ mod tests {
                 let _ = tx.send(end);
             })),
             None,
+            StreamLimits::default(),
         );
         (body, rx)
     }
@@ -520,9 +622,8 @@ mod tests {
         let (body, completion) = body_with_completion(chunks);
         let _ = body.collect().await.unwrap();
         let end = stream_end(completion).await;
-        assert!(end.transport_ok);
+        assert_eq!(end.reason, StreamEndReason::Completed);
         assert!(end.saw_error_event);
-        assert!(!end.client_disconnect);
     }
 
     #[tokio::test]
@@ -536,7 +637,7 @@ mod tests {
         let (body, completion) = body_with_completion(chunks);
         let _ = body.collect().await;
         let end = stream_end(completion).await;
-        assert!(!end.transport_ok);
+        assert_eq!(end.reason, StreamEndReason::UpstreamError);
         assert!(end.saw_error_event);
     }
 
@@ -546,9 +647,8 @@ mod tests {
         let (body, completion) = body_with_completion(chunks);
         let _ = body.collect().await;
         let end = stream_end(completion).await;
-        assert!(!end.transport_ok);
+        assert_eq!(end.reason, StreamEndReason::UpstreamError);
         assert!(!end.saw_error_event);
-        assert!(!end.client_disconnect);
     }
 
     #[tokio::test]
@@ -559,9 +659,8 @@ mod tests {
         let (body, completion) = body_with_completion(chunks);
         let _ = body.collect().await.unwrap();
         let end = stream_end(completion).await;
-        assert!(end.transport_ok);
+        assert_eq!(end.reason, StreamEndReason::Completed);
         assert!(!end.saw_error_event);
-        assert!(!end.client_disconnect);
     }
 
     #[tokio::test]
@@ -576,7 +675,6 @@ mod tests {
         let _ = stream.next().await;
         drop(stream);
         let end = stream_end(completion).await;
-        assert!(end.transport_ok);
-        assert!(end.client_disconnect);
+        assert_eq!(end.reason, StreamEndReason::ClientDisconnect);
     }
 }

experimental/sgl-router/src/server/metrics.rs

@@ -92,7 +92,7 @@
 //! The exposition is text/plain; version=0.0.4 per the Prometheus spec.
 
 use crate::config::PolicyKind;
-use crate::proxy::sse::StreamEnd;
+use crate::proxy::sse::{StreamEnd, StreamEndReason};
 use parking_lot::Mutex;
 use std::collections::HashMap;
 use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
@@ -236,14 +236,19 @@ pub enum StreamOutcome {
     UpstreamError,
     /// The client disconnected before the stream finished.
     ClientDisconnect,
+    /// The router's stale-request deadline aborted the stream.
+    Expired,
 }
 
 pub(crate) fn classify_stream_end(end: StreamEnd) -> StreamOutcome {
-    match (end.transport_ok, end.saw_error_event, end.client_disconnect) {
-        (false, _, _) => StreamOutcome::UpstreamError,
-        (_, true, _) => StreamOutcome::StreamErrorEvent,
-        (_, _, true) => StreamOutcome::ClientDisconnect,
-        _ => StreamOutcome::Ok,
+    match end.reason {
+        StreamEndReason::Expired => StreamOutcome::Expired,
+        StreamEndReason::UpstreamError
+        | StreamEndReason::IdleTimeout
+        | StreamEndReason::PumpPanicked => StreamOutcome::UpstreamError,
+        _ if end.saw_error_event => StreamOutcome::StreamErrorEvent,
+        StreamEndReason::ClientDisconnect => StreamOutcome::ClientDisconnect,
+        StreamEndReason::Completed => StreamOutcome::Ok,
     }
 }
 
@@ -254,6 +259,7 @@ impl StreamOutcome {
             Self::StreamErrorEvent => "stream_error_event",
             Self::UpstreamError => "upstream_error",
             Self::ClientDisconnect => "client_disconnect",
+            Self::Expired => "expired",
         }
     }
 }
@@ -1500,24 +1506,28 @@ mod tests {
     fn stream_outcome_precedence() {
         use StreamOutcome::*;
 
-        for (transport_ok, saw_error_event, client_disconnect, expected) in [
-            (false, false, false, UpstreamError),
-            (false, false, true, UpstreamError),
-            (false, true, false, UpstreamError),
-            (false, true, true, UpstreamError),
-            (true, false, false, Ok),
-            (true, false, true, ClientDisconnect),
-            (true, true, false, StreamErrorEvent),
-            (true, true, true, StreamErrorEvent),
+        for (reason, expected) in [
+            (StreamEndReason::Completed, Ok),
+            (StreamEndReason::ClientDisconnect, ClientDisconnect),
+            (StreamEndReason::UpstreamError, UpstreamError),
+            (StreamEndReason::IdleTimeout, UpstreamError),
+            (StreamEndReason::PumpPanicked, UpstreamError),
+            (StreamEndReason::Expired, Expired),
         ] {
+            for saw_error_event in [false, true] {
                 let end = StreamEnd {
-                transport_ok,
+                    reason,
                     saw_error_event,
-                client_disconnect,
+                };
+                let expected = if saw_error_event && matches!(expected, Ok | ClientDisconnect) {
+                    StreamErrorEvent
+                } else {
+                    expected
                 };
                 assert_eq!(classify_stream_end(end), expected, "{end:?}");
             }
         }
+    }
 
     #[test]
     fn record_stream_outcome_emits_labelled_counter_lines() {
@@ -1527,12 +1537,14 @@ mod tests {
         reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::StreamErrorEvent);
         reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::UpstreamError);
         reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::ClientDisconnect);
+        reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::Expired);
         let out = reg.render();
         for expected in [
             r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="ok"} 2"#,
             r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="stream_error_event"} 1"#,
             r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="upstream_error"} 1"#,
             r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="client_disconnect"} 1"#,
+            r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="expired"} 1"#,
         ] {
             assert_metric_line(&out, expected);
         }

experimental/sgl-router/src/server/routes/chat/forward.rs

@@ -5,7 +5,7 @@
 
 use super::preparation::{generate_room_id, BootstrapFields, PreparedChatRequest};
 use crate::discovery::WorkerMode;
-use crate::proxy::sse::StreamEnd;
+use crate::proxy::sse::{StreamEnd, StreamEndReason};
 use crate::server::app_context::AppContext;
 use crate::server::error::ApiError;
 use crate::server::metrics::{
@@ -20,6 +20,7 @@ use axum::response::IntoResponse;
 use bytes::Bytes;
 use std::sync::Arc;
 use std::time::Instant;
+use tokio_util::sync::CancellationToken;
 
 const CHAT_PATH: &str = "/v1/chat/completions";
 // Expose the selected decode worker to both PD workers and the client.
@@ -64,8 +65,6 @@ pub(super) async fn forward_chat_request(
         request.input_token_count,
         0,
     );
-    // PD requests keep using the prefill expiration token after dispatching decode.
-    let expiration_token = active_request_guard.cancel_token().clone();
     // Attribute the outcome to the worker supplying the client-visible response.
     let metrics = DispatchMetrics::new(
         ctx,
@@ -105,6 +104,9 @@ pub(super) async fn forward_chat_request(
         (prefill, prefill_load_guards)
     };
 
+    // In PD mode, prefill can finish before decode. Watch the registration
+    // held by the response so expiration remains live for its full lifetime.
+    let expiration_token = response_load_guards.1.cancel_token().clone();
     let response_future = forward_to_response_worker(
         ctx,
         &response_worker,
@@ -112,6 +114,7 @@ pub(super) async fn forward_chat_request(
         body,
         response_load_guards,
         &metrics,
+        expiration_token.clone(),
     );
     // A ready response wins if request expiration fires in the same poll.
     let result = tokio::select! {
@@ -192,6 +195,7 @@ async fn forward_to_response_worker(
     body: Bytes,
     load_guards: LoadGuards,
     metrics: &DispatchMetrics,
+    expiration: CancellationToken,
 ) -> Result<Response<Body>, ApiError> {
     if metrics.streaming {
         // Load and duration guards live until the SSE pump ends, not just until headers arrive.
@@ -208,6 +212,7 @@ async fn forward_to_response_worker(
                 Some(stream_guards),
                 Some(metrics.first_byte_callback()),
                 Some(metrics.stream_end_callback(worker.url.clone())),
+                Some(expiration),
             )
             .await
     } else {
@@ -279,6 +284,9 @@ impl DispatchMetrics {
         let metrics = Arc::clone(&self.registry);
         let model = self.model.clone();
         Box::new(move |end| {
+            if end.reason == StreamEndReason::Expired {
+                metrics.record_stale_request(StaleRequestOutcome::Expired);
+            }
             metrics.record_stream_outcome(&response_worker_url, &model, classify_stream_end(end));
         })
     }

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

@@ -1129,6 +1129,7 @@ async fn forward_streaming_to_records_failure_on_mid_stream_drop() {
             None,
             None,
             None,
+            None,
         )
         .await;
 

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

@@ -168,6 +168,7 @@ async fn h2c_client_streams_sse_from_http2_only_worker() {
                 flag.store(true, std::sync::atomic::Ordering::SeqCst);
             })),
             None,
+            None,
         )
         .await
         .expect("h2c client must stream from an HTTP/2-only worker");

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

@@ -92,6 +92,112 @@ fn chat_request() -> Request<Body> {
         .unwrap()
 }
 
+#[tokio::test]
+async fn pd_decode_stream_expires_after_prefill_completes() {
+    use sgl_router::state::load_monitor::router_inflight_load::{
+        MockClock, RouterInflightLoadRegistry,
+    };
+    use std::time::Instant;
+
+    let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await;
+    let decode = crate::common::mock_worker::MockWorker::start_slow_stream(
+        vec!["data: chunk\n\n"; 1000],
+        Duration::from_millis(10),
+    )
+    .await;
+    let cfg = config();
+    let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
+    let registry = Arc::new(WorkerRegistry::default());
+    registry
+        .add(WorkerSpec {
+            id: WorkerId("p1".into()),
+            url: prefill.url.clone(),
+            mode: WorkerMode::Prefill,
+            model_ids: vec![ModelId("tiny".into())],
+            bootstrap_port: Some(8997),
+        })
+        .unwrap();
+    registry
+        .add(WorkerSpec {
+            id: WorkerId("d1".into()),
+            url: decode.url.clone(),
+            mode: WorkerMode::Decode,
+            model_ids: vec![ModelId("tiny".into())],
+            bootstrap_port: None,
+        })
+        .unwrap();
+    let prefill_worker = registry.get(&WorkerId("p1".into())).unwrap();
+    let decode_worker = registry.get(&WorkerId("d1".into())).unwrap();
+    let clock = Arc::new(MockClock::new(Instant::now()));
+    let inflight = RouterInflightLoadRegistry::new(clock.clone(), Duration::from_secs(10));
+    let policies = Arc::new(build_registry_with_defaults(&cfg).unwrap());
+    let proxy = Arc::new(Proxy::new(Duration::from_secs(5)).unwrap());
+    let ctx = Arc::new(AppContext::with_router_inflight_load(
+        cfg,
+        tokenizers,
+        proxy,
+        registry,
+        policies,
+        inflight.clone(),
+    ));
+    let request = Request::builder()
+        .method("POST")
+        .uri("/v1/chat/completions")
+        .header("content-type", "application/json")
+        .body(Body::from(
+            serde_json::json!({
+                "model": "tiny",
+                "messages": [{"role": "user", "content": "hi"}],
+                "stream": true,
+            })
+            .to_string(),
+        ))
+        .unwrap();
+    // Two prior faults make any accidental expiry failure trip the default breaker.
+    decode_worker.breaker.record_failure();
+    decode_worker.breaker.record_failure();
+    let response = build_router(ctx.clone()).oneshot(request).await.unwrap();
+    assert_eq!(response.status(), StatusCode::OK);
+    let mut body = response.into_body();
+    tokio::time::timeout(Duration::from_secs(2), body.frame())
+        .await
+        .unwrap()
+        .unwrap()
+        .unwrap();
+    tokio::time::timeout(Duration::from_secs(2), async {
+        while inflight.inflight_count() != 1 || prefill_worker.router_inflight_load() != 0 {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .expect("prefill must complete before expiring decode");
+    assert_eq!(decode_worker.router_inflight_load(), 1);
+    clock.advance(Duration::from_secs(11));
+    assert_eq!(inflight.sweep_stale(), 1);
+    let error = tokio::time::timeout(Duration::from_secs(2), body.collect())
+        .await
+        .expect("decode stream must stop when its registration expires")
+        .unwrap_err();
+    assert!(error.to_string().contains("stale_request_timeout"));
+    assert_eq!(decode_worker.router_inflight_load(), 0);
+    assert_eq!(decode_worker.breaker.snapshot().state_code, 0);
+    let metrics = ctx.metrics.render();
+    assert!(metrics
+        .lines()
+        .any(|line| line == r#"sgl_router_stale_requests_total{outcome="expired"} 1"#));
+    let expected = format!(
+        r#"sgl_router_stream_outcome_total{{worker_url="{}",model_id="tiny",outcome="expired"}} 1"#,
+        decode.url,
+    );
+    assert!(metrics.lines().any(|line| line == expected));
+    assert!(!metrics
+        .lines()
+        .any(|line| line.starts_with("sgl_router_stream_outcome_total{")
+            && line.contains(r#"outcome="upstream_error""#)));
+    decode_worker.breaker.record_failure();
+    assert_eq!(decode_worker.breaker.snapshot().state_code, 1);
+}
+
 /// Gap closer #1: PD mode with only decode workers → 503 with
 /// `no_prefill_workers_available`. The chat route is a prefill
 /// dispatch, so a decode-only pool means partial failure.