[Router] Abort the engine when a client disconnects mid-request (#39461)

Co-authored-by: Kangyan Zhou <kangyan.zhou@radixark.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
Co-authored-by: Kan Wu <wukanustc@gmail.com>

experimental/sgl-router/tests/proxy/common/mock_worker.rs

@@ -39,9 +39,22 @@ pub struct MockWorker {
     // Used in header_forwarding_test; not every test file reads captured headers.
     #[allow(dead_code)]
     pub captured: Arc<Mutex<CapturedHeaders>>,
+    #[allow(dead_code)]
+    pub abort_log: Arc<Mutex<Vec<Value>>>,
     _shutdown: oneshot::Sender<()>,
 }
 
+#[allow(dead_code)] // shared across all axum variants
+fn abort_request_route<S>(log: Arc<Mutex<Vec<Value>>>) -> axum::routing::MethodRouter<S>
+where
+    S: Clone + Send + Sync + 'static,
+{
+    post(move |Json(body): Json<Value>| async move {
+        log.lock().unwrap().push(body);
+        StatusCode::OK
+    })
+}
+
 impl MockWorker {
     /// Bind to a random port on 127.0.0.1 and start serving.
     ///
@@ -50,6 +63,7 @@ impl MockWorker {
     #[allow(dead_code)] // Only used by some test files.
     pub async fn start(stream_chunks: Vec<&'static str>) -> Self {
         let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
+        let abort_log: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
         let state = MockWorkerState {
             captured: captured.clone(),
             stream_chunks: Arc::new(stream_chunks),
@@ -60,6 +74,7 @@ impl MockWorker {
         let app = axum::Router::new()
             .route("/v1/chat/completions", post(chat))
             .route("/server_info", get(serve_tiny_server_info))
+            .route("/abort_request", abort_request_route(abort_log.clone()))
             .with_state(state);
 
         let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -77,6 +92,7 @@ impl MockWorker {
         Self {
             url,
             captured,
+            abort_log,
             _shutdown: tx,
         }
     }
@@ -88,6 +104,7 @@ impl MockWorker {
     #[allow(dead_code)]
     pub async fn start_hanging(delay: Duration) -> Self {
         let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
+        let abort_log: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
 
         #[derive(Clone)]
         struct HangState {
@@ -127,6 +144,7 @@ impl MockWorker {
         let app = axum::Router::new()
             .route("/v1/chat/completions", post(hang_handler))
             .route("/server_info", get(serve_tiny_server_info))
+            .route("/abort_request", abort_request_route(abort_log.clone()))
             .with_state(state);
 
         let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -144,6 +162,7 @@ impl MockWorker {
         Self {
             url,
             captured,
+            abort_log,
             _shutdown: tx,
         }
     }
@@ -154,6 +173,7 @@ impl MockWorker {
     #[allow(dead_code)]
     pub async fn start_slow_stream(chunks: Vec<&'static str>, delay: Duration) -> Self {
         let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
+        let abort_log: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
 
         #[derive(Clone)]
         struct SlowState {
@@ -207,6 +227,7 @@ impl MockWorker {
         let app = axum::Router::new()
             .route("/v1/chat/completions", post(slow_chat))
             .route("/server_info", get(serve_tiny_server_info))
+            .route("/abort_request", abort_request_route(abort_log.clone()))
             .with_state(state);
 
         let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -224,6 +245,7 @@ impl MockWorker {
         Self {
             url,
             captured,
+            abort_log,
             _shutdown: tx,
         }
     }
@@ -248,6 +270,7 @@ impl MockWorker {
         partial_body_bytes: &'static [u8],
     ) -> Self {
         let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
+        let abort_log: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
         let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
         let addr: SocketAddr = listener.local_addr().unwrap();
         let url = format!("http://{addr}");
@@ -309,6 +332,7 @@ impl MockWorker {
         Self {
             url,
             captured,
+            abort_log,
             _shutdown: tx,
         }
     }
@@ -319,6 +343,7 @@ impl MockWorker {
     #[allow(dead_code)]
     pub async fn start_returning_error(status: StatusCode, body: Value) -> Self {
         let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
+        let abort_log: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
         let body_arc = Arc::new(body.to_string());
 
         #[derive(Clone)]
@@ -360,6 +385,7 @@ impl MockWorker {
         let app = axum::Router::new()
             .route("/v1/chat/completions", post(error_handler))
             .route("/server_info", get(serve_tiny_server_info))
+            .route("/abort_request", abort_request_route(abort_log.clone()))
             .with_state(state);
 
         let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -377,6 +403,7 @@ impl MockWorker {
         Self {
             url,
             captured,
+            abort_log,
             _shutdown: tx,
         }
     }