[sgl-router] Fail fast and cancel decode when prefill fails

experimental/sgl-router/monitoring/README.md

@@ -26,7 +26,7 @@ Families the router emits. The dashboard graphs all of them except the
 | `sgl_router_worker_requests_total` | Counter | Per-worker **dispatches** by `worker_url`, `model_id`, `mode`, `outcome` (recorded after dispatch; blind to pre-dispatch drops). See [Dispatch outcomes](#dispatch-outcomes) |
 | `sgl_router_request_duration_seconds` | Histogram | End-to-end request latency by `model_id` |
 | `sgl_router_ttft_seconds` | Histogram | Time to first token (streaming) by `model_id` |
-| `sgl_router_stream_outcome_total` | Counter | Streaming outcomes by `worker_url`, `model_id`, and `outcome` (`ok`, `stream_error_event`, `upstream_error`, or `client_disconnect`). Counts committed 2xx streams only — non-2xx responses are counted by status in `responses_total` |
+| `sgl_router_stream_outcome_total` | Counter | Streaming outcomes by `worker_url`, `model_id`, and `outcome` (`ok`, `stream_error_event`, `upstream_error`, `client_disconnect`, `expired`, or `aborted`). Counts committed 2xx streams only — non-2xx responses are counted by status in `responses_total` |
 | `sgl_router_active_load` | Gauge | Per-worker prefill-token / decode-block load |
 | `sgl_router_workers` | Gauge | Registered worker count by `mode` |
 | `sgl_router_worker_health` | Gauge | Per-worker health (1=breaker admits, 0=open) |
@@ -126,7 +126,7 @@ worker error the router forwards is a successful *proxy* operation and a failed
 | `client_error` | 4xx except 429 | no — the caller sent something invalid |
 | `backpressure` | 429, 503 | no — responsive but at capacity |
 | `error` | 5xx except 503, plus transport failures, timeouts and incomplete bodies | **yes** |
-| `cancelled` | the router's own stale-request deadline | no |
+| `cancelled` | the router's own stale-request deadline, or a PD decode dispatch dropped because prefill failed first | no |
 
 `error` is the only bucket that means *this worker failed*, which is why the
 Error-ratio panel uses it alone. The split matters during an incident: a
@@ -136,8 +136,11 @@ statuses — so the two agree, and the error ratio keeps pointing at genuine
 faults instead of pegging at 100% exactly when it is being read.
 
 A hung worker surfaces as `error` (the router's upstream timeout), *not* as
-`cancelled`. Only the stale-request deadline produces `cancelled`;
-`sgl_router_stale_requests_total{outcome="expired"}` counts the same events.
+`cancelled`. Two things produce `cancelled`: the stale-request deadline, which
+`sgl_router_stale_requests_total{outcome="expired"}` counts the same events of,
+and a PD decode worker whose dispatch was dropped because prefill failed first —
+the client's status is then prefill's, booked against the prefill worker, so
+decode's abandoned dispatch is counted here rather than disappearing.
 
 ## Access log
 

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

@@ -49,7 +49,7 @@ enum BreakerOutcome {
     /// A real fault (5xx other than backpressure) → `record_failure`: count
     /// toward opening.
     Failure,
-    /// Backpressure or router-side stream expiry →
+    /// Backpressure or a router-side stream expiry or abort →
     /// `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.
@@ -85,13 +85,13 @@ fn breaker_outcome(status: reqwest::StatusCode) -> BreakerOutcome {
     }
 }
 
-/// Router-side expiry says nothing about worker health. Preserve the existing
+/// Router-side expiry or abort 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::Expired | StreamEndReason::Aborted => BreakerOutcome::Neutral,
         StreamEndReason::Completed | StreamEndReason::ClientDisconnect => BreakerOutcome::Success,
         StreamEndReason::UpstreamError
         | StreamEndReason::IdleTimeout
@@ -314,6 +314,7 @@ impl Proxy {
         on_first_byte: Option<Box<dyn FnOnce() + Send + 'static>>,
         on_stream_end: Option<Box<dyn FnOnce(sse::StreamEnd) + Send + 'static>>,
         expiration: Option<CancellationToken>,
+        stream_abort: Option<CancellationToken>,
     ) -> Result<Response<Body>, ApiError> {
         let permit = breaker.acquire().ok_or_else(|| ApiError::BreakerOpen {
             worker: worker_url.to_string(),
@@ -410,6 +411,7 @@ impl Proxy {
             sse::StreamLimits {
                 idle_timeout: self.stream_idle_timeout,
                 expiration,
+                abort: stream_abort.filter(|_| status.is_success()),
             },
         );
         let mut out = Response::new(body);
@@ -500,6 +502,7 @@ mod tests {
                 None,
                 None,
                 None,
+                None,
             )
             .now_or_never()
             .is_none());
@@ -604,6 +607,7 @@ mod tests {
                 None,
                 None,
                 expiration,
+                None,
             )
             .await
             .unwrap();
@@ -843,6 +847,7 @@ mod tests {
                     None,
                     None,
                     None,
+                    None,
                 )
                 .await
                 .expect("streaming dispatch should reach the worker");

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

@@ -4,6 +4,7 @@
 //! SSE passthrough — bridges a reqwest `bytes_stream()` into an axum Body.
 
 use std::panic::AssertUnwindSafe;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::time::Duration;
 
 use axum::body::Body;
@@ -20,6 +21,8 @@ pub enum StreamEndReason {
     IdleTimeout,
     /// The router's stale-request deadline expired, regardless of worker health.
     Expired,
+    /// The router aborted the stream before its first chunk.
+    Aborted,
     ClientDisconnect,
     PumpPanicked,
 }
@@ -69,21 +72,29 @@ impl ErrorEventScanner {
     }
 }
 
+/// Whether a buffered SSE body carries an error event.
+pub fn has_error_event(body: &[u8]) -> bool {
+    ErrorEventScanner::default().feed(body)
+}
+
 /// 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>,
+    /// Aborts the stream if it fires before the first upstream chunk.
+    pub abort: Option<CancellationToken>,
 }
 
 /// Bridge a byte stream into an axum Body that streams chunks unchanged.
 ///
 /// 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.
+/// elapses between chunks, when `limits.expiration` fires, or when
+/// `limits.abort` fires before the first chunk.
 ///
 /// 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
@@ -111,10 +122,12 @@ where
         };
         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,
+        let started = AtomicBool::new(false);
+        let expired = cancelled(limits.expiration);
+        let aborted = async {
+            cancelled(limits.abort).await;
+            if started.load(Ordering::Relaxed) {
+                std::future::pending::<()>().await;
             }
         };
         // Disconnect and expiration race the whole forwarding loop, so they
@@ -130,6 +143,10 @@ where
                     end.reason = StreamEndReason::Expired;
                     Err(std::io::Error::other("SSE stream exceeded stale_request_timeout"))
                 }
+                _ = aborted => {
+                    end.reason = StreamEndReason::Aborted;
+                    Err(std::io::Error::other("SSE stream aborted by the router"))
+                }
                 result = async {
                     loop {
                         let bytes = match tokio::time::timeout(idle, stream.next()).await {
@@ -144,6 +161,7 @@ where
                                 Err(std::io::Error::other("SSE upstream idle timeout")),
                             ),
                         };
+                        started.store(true, Ordering::Relaxed);
                         if let Some(hook) = on_first_byte.take() {
                             hook();
                         }
@@ -191,6 +209,13 @@ where
     Body::from_stream(ReceiverStream::new(rx).map(Ok).chain(terminal))
 }
 
+async fn cancelled(token: Option<CancellationToken>) {
+    match token {
+        Some(token) => token.cancelled().await,
+        None => std::future::pending().await,
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -249,6 +274,7 @@ mod tests {
             StreamLimits {
                 idle_timeout: Some(Duration::from_secs(1)),
                 expiration: None,
+                abort: None,
             },
         );
         assert!(body
@@ -282,6 +308,7 @@ mod tests {
             StreamLimits {
                 idle_timeout: None,
                 expiration: Some(token.clone()),
+                abort: None,
             },
         );
         tokio::task::yield_now().await;
@@ -300,6 +327,39 @@ mod tests {
             .to_string()
             .contains("stale_request_timeout"));
     }
+
+    #[tokio::test]
+    async fn abort_ends_the_stream_only_before_its_first_chunk() {
+        for started in [false, true] {
+            let token = CancellationToken::new();
+            let (tx, rx) = tokio::sync::mpsc::channel(1);
+            let (body, end) = limited_body(
+                ReceiverStream::new(rx),
+                StreamLimits {
+                    abort: Some(token.clone()),
+                    ..Default::default()
+                },
+            );
+            let mut data = body.into_data_stream();
+            if started {
+                tx.send(Ok(Bytes::from_static(b"data: a\n\n")))
+                    .await
+                    .unwrap();
+                data.next().await.unwrap().unwrap();
+            }
+            token.cancel();
+            drop(tx);
+            let rest: Vec<_> = StreamExt::collect(data).await;
+            assert_eq!(rest.iter().any(Result::is_err), !started);
+            let expected = if started {
+                StreamEndReason::Completed
+            } else {
+                StreamEndReason::Aborted
+            };
+            assert_eq!(end.await.unwrap().reason, expected);
+        }
+    }
+
     #[tokio::test]
     async fn passes_through_a_simple_byte_stream() {
         let chunks = vec![

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

@@ -9,13 +9,8 @@ use thiserror::Error;
 
 pub const X_ROUTER_ERROR_CODE: HeaderName = HeaderName::from_static("x-router-error-code");
 
-/// Carries the worker's *own* HTTP status when the router had to synthesize a
-/// status of its own over a worker that did respond (today: a mid-body drop,
-/// where headers arrived but the body did not). Lets a gateway / operator
-/// recover what the engine actually said instead of seeing only the router's
-/// synthesized 502. Absent on every other response: a forwarded worker response
-/// already carries the worker's status in the status line, and a router-only
-/// condition (no workers, breaker open, ...) has no upstream status to report.
+/// Original worker status when a failed prefill or incomplete body becomes a router 502.
+/// Absent when no upstream response arrived or its status was forwarded unchanged.
 pub const X_ROUTER_UPSTREAM_STATUS: HeaderName =
     HeaderName::from_static("x-router-upstream-status");
 
@@ -103,6 +98,9 @@ pub enum ApiError {
     #[error("upstream response body incomplete after status {status}")]
     UpstreamStatus { status: StatusCode },
 
+    #[error("prefill failed before KV transfer completed")]
+    PrefillFailed { status: Option<StatusCode> },
+
     /// Wall-clock timeout exceeded while waiting for the upstream worker's
     /// response (per-request `request_timeout`).
     ///
@@ -187,7 +185,9 @@ impl ApiError {
             ApiError::SamplingContract { .. } => ErrorClass::BadRequest,
             ApiError::ModelNotFound(_) => ErrorClass::NotFound,
             ApiError::UpstreamUnreachable { .. } => ErrorClass::Upstream,
-            ApiError::UpstreamStatus { .. } => ErrorClass::Upstream,
+            ApiError::UpstreamStatus { .. } | ApiError::PrefillFailed { .. } => {
+                ErrorClass::Upstream
+            }
             ApiError::UpstreamTimeout { .. } => ErrorClass::Timeout,
             ApiError::NoHealthyWorkers { .. } => ErrorClass::NoTarget,
             ApiError::NoPrefillWorkersAvailable { .. } => ErrorClass::NoTarget,
@@ -218,6 +218,7 @@ impl ApiError {
             // didn't. We surface a 502 but echo the worker's status in
             // `x-router-upstream-status` (see `into_response`).
             ApiError::UpstreamStatus { .. } => "upstream_body_incomplete",
+            ApiError::PrefillFailed { .. } => "prefill_failed",
             ApiError::UpstreamTimeout { .. } => "upstream_timeout",
             ApiError::NoHealthyWorkers { .. } => "no_healthy_workers",
             ApiError::NoPrefillWorkersAvailable { .. } => "no_prefill_workers_available",
@@ -232,13 +233,14 @@ impl ApiError {
 
     /// The worker's *own* status to echo in `x-router-upstream-status`, for the
     /// case where the router synthesized its own status over a worker that did
-    /// respond. Today only the mid-body-drop (`UpstreamStatus`) carries one. This
+    /// respond, including a failed prefill or an incomplete response body. This
     /// is an exhaustive, wildcard-free match (not an `if let` at the call site) so
     /// a future "synthesized over a responding worker" variant is forced to decide
     /// whether it echoes a status, rather than silently inheriting `None`.
     fn upstream_status(&self) -> Option<StatusCode> {
         match self {
             ApiError::UpstreamStatus { status } => Some(*status),
+            ApiError::PrefillFailed { status } => *status,
             ApiError::BadRequest(_)
             | ApiError::SamplingContract { .. }
             | ApiError::ModelNotFound(_)
@@ -311,6 +313,7 @@ impl IntoResponse for ApiError {
                 );
                 "upstream response was incomplete".to_string()
             }
+            ApiError::PrefillFailed { .. } => self.to_string(),
             ApiError::UpstreamTimeout { worker } => {
                 tracing::warn!(upstream = %worker, "upstream request timed out");
                 "upstream request timed out".to_string()
@@ -369,9 +372,7 @@ impl IntoResponse for ApiError {
             .into_response();
         resp.headers_mut()
             .insert(X_ROUTER_ERROR_CODE, HeaderValue::from_static(code));
-        // Preserve the worker's real status when we synthesized our own (today,
-        // only the mid-body-drop case: the worker sent a status, then dropped the
-        // body, so we report a 502 but don't throw away what it said).
+        // Preserve the worker's real status when we synthesized our own.
         if let Some(upstream) = self.upstream_status() {
             resp.headers_mut().insert(
                 X_ROUTER_UPSTREAM_STATUS,

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

@@ -179,8 +179,9 @@ pub enum RequestOutcome {
     /// The worker failed to serve the request: a 5xx fault, a transport failure,
     /// a timeout, or a body that never completed.
     Error,
-    /// The router cancelled the request itself — today only the stale-request
-    /// deadline. Never derived from a status; see [`outcome_from_status`].
+    /// The router cancelled the request itself: the stale-request deadline, or
+    /// a PD decode dispatch abandoned because prefill failed first. Never
+    /// derived from a status; see [`outcome_from_status`].
     Cancelled,
 }
 
@@ -260,11 +261,14 @@ pub enum StreamOutcome {
     ClientDisconnect,
     /// The router's stale-request deadline aborted the stream.
     Expired,
+    /// The router aborted the stream before its first chunk.
+    Aborted,
 }
 
 pub(crate) fn classify_stream_end(end: StreamEnd) -> StreamOutcome {
     match end.reason {
         StreamEndReason::Expired => StreamOutcome::Expired,
+        StreamEndReason::Aborted => StreamOutcome::Aborted,
         StreamEndReason::UpstreamError
         | StreamEndReason::IdleTimeout
         | StreamEndReason::PumpPanicked => StreamOutcome::UpstreamError,
@@ -282,6 +286,7 @@ impl StreamOutcome {
             Self::UpstreamError => "upstream_error",
             Self::ClientDisconnect => "client_disconnect",
             Self::Expired => "expired",
+            Self::Aborted => "aborted",
         }
     }
 }
@@ -1602,6 +1607,7 @@ mod tests {
             (StreamEndReason::IdleTimeout, UpstreamError),
             (StreamEndReason::PumpPanicked, UpstreamError),
             (StreamEndReason::Expired, Expired),
+            (StreamEndReason::Aborted, Aborted),
         ] {
             for saw_error_event in [false, true] {
                 let end = StreamEnd {