minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Kangyan-Zhou
2026-09-22 12:35:56 +08:00
committed by GitHub
co-authored by Kangyan Zhou Claude Opus 5 Shangming Cai Kan Wu
parent a9f02b0fa4
commit 2032f3a071
15 changed files with 629 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs
+12 -1
View File
@@ -198,7 +198,18 @@ async fn caller_input_ids_are_used_for_routing_and_preserved() {
send(Arc::clone(&ctx), request.clone()).await,
StatusCode::OK
);
assert_eq!(captured(&mock), request, "body must be forwarded untouched");
let mut forwarded = captured(&mock);
let rid = forwarded
.as_object_mut()
.expect("a forwarded chat body is an object")
.remove("rid");
assert!(
rid.as_ref()
.and_then(Value::as_str)
.is_some_and(crate::common::is_engine_shaped_rid),
"plain mode must mint an abort rid; got {rid:?}",
);
assert_eq!(forwarded, request, "body must be forwarded untouched");
}
// Bypasses are not rendering failures.
assert!(!ctx
experimental/sgl-router/tests/proxy/chat_routing.rs
+93 -47
View File
@@ -17,10 +17,14 @@ use sgl_router::workers::{WireProtocol, Worker, WorkerRegistry};
use axum::body::Body;
use axum::http::{Request, StatusCode};
use http_body_util::BodyExt;
use sgl_router::state::load_monitor::router_inflight_load::{
spawn_janitor, JanitorHandle, RouterInflightLoadRegistry,
};
use std::sync::Arc;
use std::time::Duration;
use tower::ServiceExt;
mod cancellation;
mod reorg;
const TEST_TIMEOUT: Duration = Duration::from_secs(5);
@@ -74,6 +78,36 @@ fn build_ctx_with_worker(url: &str) -> Arc<AppContext> {
Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies))
}
/// Expire requests after 50ms; keep the janitor handle alive during the test.
fn build_ctx_with_janitor(url: &str) -> (Arc<AppContext>, JanitorHandle) {
let cfg = config_for(url);
let registry = Arc::new(WorkerRegistry::default());
let _ = registry.add(WorkerSpec {
id: WorkerId("w1".into()),
url: url.to_string(),
mode: WorkerMode::Plain,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
});
let policies = Arc::new(build_policy_registry(&cfg).unwrap());
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let proxy = Arc::new(Proxy::new(TEST_TIMEOUT).unwrap());
let router_inflight_load = RouterInflightLoadRegistry::new(
Arc::new(sgl_router::state::load_monitor::router_inflight_load::SystemTimeClock),
Duration::from_millis(50),
);
let janitor = spawn_janitor(Arc::clone(&router_inflight_load), Duration::from_millis(20));
let ctx = Arc::new(AppContext::with_router_inflight_load(
cfg,
tokenizers,
proxy,
registry,
policies,
router_inflight_load,
));
(ctx, janitor)
}
#[tokio::test]
async fn non_streaming_returns_200() {
let worker = crate::common::mock_worker::MockWorker::start(vec![]).await;
@@ -1033,6 +1067,7 @@ async fn forward_json_to_records_failure_on_body_drop() {
"/v1/chat/completions",
&headers,
body,
None,
)
.await;
assert!(res.is_err(), "body drop should surface as ApiError");
@@ -1090,6 +1125,7 @@ async fn forward_json_to_records_success_only_after_body_completes() {
"/v1/chat/completions",
&headers,
bytes::Bytes::from_static(b"{}"),
None,
)
.await;
assert!(res.is_ok(), "clean OK call must succeed: {res:?}");
@@ -1145,6 +1181,7 @@ async fn forward_streaming_to_records_failure_on_mid_stream_drop() {
None,
None,
None,
None,
)
.await;
@@ -1248,6 +1285,7 @@ async fn forward_json_to_records_failure_on_5xx() {
"/v1/chat/completions",
&headers,
body,
None,
)
.await;
@@ -1282,6 +1320,7 @@ async fn forward_json_to_rejects_when_breaker_open() {
"/v1/chat/completions",
&headers,
body,
None,
)
.await;
@@ -1320,6 +1359,7 @@ async fn forward_json_to_malformed_url_returns_worker_misconfigured_and_trips_br
"/v1/chat/completions",
&headers,
body,
None,
)
.await;
@@ -1607,58 +1647,17 @@ async fn streaming_active_load_drops_on_client_disconnect() {
/// returns; cancellation fires; handler returns 504.
#[tokio::test]
async fn janitor_expiry_returns_504_stale_request_expired() {
use sgl_router::state::load_monitor::router_inflight_load::{
spawn_janitor, RouterInflightLoadRegistry,
};
// Upstream that takes 2s to respond — longer than our 50ms
// stale_request_timeout.
// Upstream that takes 2s to respond — longer than the helper's 50ms
// stale_request_timeout, so the janitor sweeps before it answers.
let worker =
crate::common::mock_worker::MockWorker::start_hanging(Duration::from_secs(2)).await;
let cfg = config_for(&worker.url);
let registry = Arc::new(WorkerRegistry::default());
let _ = registry.add(WorkerSpec {
id: WorkerId("w1".into()),
url: worker.url.clone(),
mode: WorkerMode::Plain,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
});
let policies = Arc::new(build_policy_registry(&cfg).unwrap());
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let proxy = Arc::new(Proxy::new(TEST_TIMEOUT).unwrap());
// Aggressive 50ms timeout: the janitor will sweep on the next
// tick (every 20ms) and fire the cancellation token before the
// upstream returns.
let router_inflight_load = RouterInflightLoadRegistry::new(
Arc::new(sgl_router::state::load_monitor::router_inflight_load::SystemTimeClock),
Duration::from_millis(50),
);
let _janitor = spawn_janitor(Arc::clone(&router_inflight_load), Duration::from_millis(20));
let ctx = Arc::new(AppContext::with_router_inflight_load(
cfg,
tokenizers,
proxy,
registry,
policies,
router_inflight_load,
));
let (ctx, _janitor) = build_ctx_with_janitor(&worker.url);
let app = build_router(ctx);
let req = Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("content-type", "application/json")
.body(Body::from(
serde_json::to_vec(&serde_json::json!({
"model": "tiny",
"messages": [{"role": "user", "content": "hi"}],
"stream": false,
}))
.unwrap(),
))
let res = app
.oneshot(cancellation::request(serde_json::json!({})))
.await
.unwrap();
let res = app.oneshot(req).await.unwrap();
assert_eq!(
res.status(),
StatusCode::GATEWAY_TIMEOUT,
@@ -1677,6 +1676,53 @@ async fn janitor_expiry_returns_504_stale_request_expired() {
body_str.contains("\"code\":\"stale_request_expired\""),
"504 body must encode the same code in the JSON envelope: {body_str}",
);
assert_engine_abort(&worker).await;
}
#[tokio::test]
async fn janitor_expiry_aborts_before_headers_and_mid_stream() {
use crate::common::mock_worker::MockWorker;
for before_headers in [true, false] {
let worker = if before_headers {
MockWorker::start_hanging(Duration::from_secs(2)).await
} else {
MockWorker::start_slow_stream(vec!["data: a\n\n"], Duration::from_secs(2)).await
};
let (ctx, _janitor) = build_ctx_with_janitor(&worker.url);
let response = build_router(ctx)
.oneshot(cancellation::request(serde_json::json!({"stream":true})))
.await
.unwrap();
assert_eq!(
response.status(),
if before_headers {
StatusCode::GATEWAY_TIMEOUT
} else {
StatusCode::OK
}
);
let result = response.into_body().collect().await;
assert_eq!(result.is_ok(), before_headers);
assert_engine_abort(&worker).await;
}
}
async fn assert_engine_abort(worker: &crate::common::mock_worker::MockWorker) {
tokio::time::timeout(TEST_TIMEOUT, async {
while worker.abort_log.lock().unwrap().is_empty() {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.unwrap();
let forwarded: serde_json::Value =
serde_json::from_slice(worker.captured.lock().unwrap().last_body.as_ref().unwrap())
.unwrap();
assert_eq!(
*worker.abort_log.lock().unwrap(),
vec![serde_json::json!({"rid":forwarded["rid"], "abort_all":false})]
);
}
/// Task A: a non-streaming request that errors out (upstream
experimental/sgl-router/tests/proxy/chat_routing/cancellation.rs
+205
View File
@@ -0,0 +1,205 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use super::*;
use axum::{extract::State, http::HeaderMap, routing::post, Json, Router};
use serde_json::{json, Value};
use tokio::sync::mpsc;
type Event = (&'static str, Value);
type Events = mpsc::UnboundedSender<Event>;
async fn chat(State(events): State<Events>, Json(body): Json<Value>) -> (StatusCode, Body) {
events.send(("chat", body.clone())).unwrap();
if body["before_headers"] == true || (body["hold"] == true && body["stream"] != true) {
std::future::pending::<()>().await;
}
let response = if body["hold"] == true {
Body::from_stream(futures::stream::pending::<
Result<bytes::Bytes, std::io::Error>,
>())
} else if body["stream"] == true {
Body::from("data: [DONE]\n\n")
} else {
Body::from("{}")
};
(
StatusCode::from_u16(body["status"].as_u64().unwrap_or(200) as u16).unwrap(),
response,
)
}
async fn abort(
State(events): State<Events>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> StatusCode {
assert_eq!(headers["authorization"], "Bearer test");
events.send(("abort", body)).unwrap();
StatusCode::INTERNAL_SERVER_ERROR // Abort failures must not affect the worker's breaker.
}
struct Harness {
ctx: Arc<AppContext>,
events: mpsc::UnboundedReceiver<Event>,
server: tokio::task::JoinHandle<()>,
}
impl Harness {
async fn new() -> Self {
let (events, rx) = mpsc::unbounded_channel();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let ctx = build_ctx_with_worker(&format!("http://{}", listener.local_addr().unwrap()));
let server = tokio::spawn(async move {
axum::serve(
listener,
Router::new()
.route("/v1/chat/completions", post(chat))
.route("/abort_request", post(abort))
.with_state(events),
)
.await
.unwrap();
});
Self {
ctx,
events: rx,
server,
}
}
async fn event(&mut self, expected: &str) -> Value {
let (kind, body) = tokio::time::timeout(TEST_TIMEOUT, self.events.recv())
.await
.unwrap()
.unwrap();
assert_eq!(kind, expected);
body
}
async fn quiet(&mut self) {
assert!(
tokio::time::timeout(Duration::from_millis(50), self.events.recv())
.await
.is_err()
);
}
}
impl Drop for Harness {
fn drop(&mut self) {
self.server.abort();
}
}
pub(super) fn request(mut body: Value) -> Request<Body> {
body["model"] = json!("tiny");
Request::post("/v1/chat/completions")
.header("content-type", "application/json")
.header("authorization", "Bearer test")
.header("x-request-id", "reused-gateway-id")
.body(Body::from(body.to_string()))
.unwrap()
}
#[tokio::test]
async fn only_unfinished_requests_abort() {
let mut h = Harness::new().await;
for (stream, hold, before_headers) in [
(false, false, false),
(true, false, false),
(false, true, false),
(true, true, true),
(true, true, false),
] {
let task = tokio::spawn(build_router(h.ctx.clone()).oneshot(request(json!({
"stream": stream, "hold": hold, "before_headers": before_headers,
}))));
let forwarded = h.event("chat").await;
assert!(crate::common::is_engine_shaped_rid(
forwarded["rid"].as_str().unwrap()
));
let worker = h.ctx.registry.get(&WorkerId("w1".into())).unwrap();
if hold && (!stream || before_headers) {
worker.breaker.record_failure();
worker.breaker.record_failure();
task.abort();
assert!(task.await.unwrap_err().is_cancelled());
} else {
let response = tokio::time::timeout(TEST_TIMEOUT, task)
.await
.unwrap()
.unwrap()
.unwrap();
if hold {
drop(response); // Silent upstream: cancellation must not wait for a token.
} else {
response.into_body().collect().await.unwrap();
}
}
if hold {
assert_eq!(
h.event("abort").await,
json!({"rid": forwarded["rid"], "abort_all": false})
);
}
h.quiet().await;
assert_eq!(worker.breaker.snapshot().state_code, 0);
worker.breaker.record_success();
}
}
#[tokio::test]
async fn caller_ids_fan_out_and_rejected_streams_do_not_abort() {
let mut h = Harness::new().await;
for fields in [
json!({"rid":"a"}),
json!({"rid":["a","b"]}),
json!({"n":2}),
json!({"status":400}),
json!({"status":429}),
json!({"status":500}),
json!({"status":503}),
] {
let mut body = json!({"stream":true, "hold":true});
body.as_object_mut()
.unwrap()
.extend(fields.as_object().unwrap().clone());
let response = build_router(h.ctx.clone())
.oneshot(request(body))
.await
.unwrap();
let forwarded = h.event("chat").await;
if fields.get("status").is_none() {
assert_eq!(forwarded.get("rid"), fields.get("rid"));
}
drop(response);
h.quiet().await;
}
}
#[tokio::test]
async fn concurrent_requests_with_the_same_header_get_distinct_abort_ids() {
let mut h = Harness::new().await;
let mut tasks = Vec::new();
for _ in 0..2 {
tasks.push(tokio::spawn(
build_router(h.ctx.clone()).oneshot(request(json!({"hold":true}))),
));
}
let mut rids = Vec::new();
for _ in 0..2 {
rids.push(h.event("chat").await["rid"].clone());
}
assert_ne!(rids[0], rids[1]);
for task in tasks {
task.abort();
assert!(task.await.unwrap_err().is_cancelled());
}
for _ in 0..2 {
let aborted = h.event("abort").await;
let index = rids.iter().position(|rid| rid == &aborted["rid"]).unwrap();
rids.remove(index);
}
h.quiet().await;
}
experimental/sgl-router/tests/proxy/chat_routing/reorg.rs
+14 -1
View File
@@ -230,7 +230,19 @@ async fn length_selects_plain_bucket_before_engine_selection() {
let response = app.clone().oneshot(request(body("hi"))).await.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let _ = response.into_body().collect().await.unwrap();
assert!(short_worker.captured.lock().unwrap().last_body.is_some());
let forwarded: serde_json::Value = serde_json::from_slice(
short_worker
.captured
.lock()
.unwrap()
.last_body
.as_ref()
.unwrap(),
)
.unwrap();
assert!(crate::common::is_engine_shaped_rid(
forwarded["rid"].as_str().unwrap()
));
assert!(long_worker.captured.lock().unwrap().last_body.is_none());
let response = app
@@ -293,6 +305,7 @@ async fn pd_picks_both_groups_from_selected_bucket_and_shares_bootstrap() {
let d: serde_json::Value =
serde_json::from_slice(decode.captured.lock().unwrap().last_body.as_ref().unwrap())
.unwrap();
assert!(p.get("rid").is_none() && d.get("rid").is_none());
assert!(p["bootstrap_room"].is_number());
assert_eq!(p["bootstrap_room"], d["bootstrap_room"]);
let calls = policy.calls.lock().unwrap();
experimental/sgl-router/tests/proxy/common/mock_worker.rs
+27
View File
@@ -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,
}
}
experimental/sgl-router/tests/proxy/common/mod.rs
+8
View File
@@ -6,3 +6,11 @@
pub mod cache_aware_fixture;
pub mod mock_worker;
pub mod streaming;
#[allow(dead_code)] // not every test file inspects forwarded rids
pub fn is_engine_shaped_rid(rid: &str) -> bool {
rid.len() == 32
&& rid
.bytes()
.all(|b| b.is_ascii_hexdigit() && !b.is_ascii_uppercase())
}
experimental/sgl-router/tests/proxy/h2c_forward.rs
+3
View File
@@ -72,6 +72,7 @@ async fn h2c_client_reaches_http2_only_worker() {
"/v1/chat/completions",
&axum::http::HeaderMap::new(),
Bytes::from_static(b"{}"),
None,
)
.await
.expect("h2c client must reach an HTTP/2-only worker");
@@ -96,6 +97,7 @@ async fn http1_client_cannot_reach_http2_only_worker() {
"/v1/chat/completions",
&axum::http::HeaderMap::new(),
Bytes::from_static(b"{}"),
None,
)
.await;
assert!(
@@ -164,6 +166,7 @@ async fn h2c_client_streams_sse_from_http2_only_worker() {
&axum::http::HeaderMap::new(),
Bytes::from_static(b"{}"),
None,
None,
Some(Box::new(move || {
flag.store(true, std::sync::atomic::Ordering::SeqCst);
})),
experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs
+58
View File
@@ -376,3 +376,61 @@ async fn pd_mode_prefill_5xx_does_not_poison_decode_response() {
let pv = parse_body(&prefill_body);
assert_eq!(bootstrap_port(&pv), Some(8997));
}
fn streaming_chat_request() -> Request<Body> {
Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("content-type", "application/json")
.body(Body::from(
serde_json::to_vec(&json!({
"model": "tiny",
"messages": [{"role": "user", "content": "hi"}],
"stream": true,
}))
.unwrap(),
))
.unwrap()
}
#[tokio::test]
async fn pd_mode_disconnect_does_not_abort_either_worker() {
let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await;
let decode = crate::common::mock_worker::MockWorker::start_slow_stream(
vec!["data: a\n\n", "data: b\n\n", "data: c\n\n"],
Duration::from_millis(50),
)
.await;
let ctx = build_ctx(vec![
WorkerSpec {
id: WorkerId("p1".into()),
url: prefill.url.clone(),
mode: WorkerMode::Prefill,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: Some(8997),
},
WorkerSpec {
id: WorkerId("d1".into()),
url: decode.url.clone(),
mode: WorkerMode::Decode,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
},
]);
let app = build_router(ctx);
let res = app.oneshot(streaming_chat_request()).await.unwrap();
assert_eq!(res.status(), StatusCode::OK);
use futures::StreamExt;
let mut data_stream = res.into_body().into_data_stream();
assert!(data_stream.next().await.is_some());
drop(data_stream);
tokio::time::sleep(Duration::from_millis(100)).await;
for worker in [&prefill, &decode] {
assert!(worker.abort_log.lock().unwrap().is_empty());
let body = await_captured_body(worker, Duration::from_secs(2), "PD worker").await;
assert!(parse_body(&body).get("rid").is_none());
}
}
experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs
+16 -2
View File
@@ -106,9 +106,23 @@ fn without_forwarding(mut cfg: Config, policy: PolicyKind) -> Config {
cfg
}
fn without_minted_rid(mut body: Value) -> Value {
let rid = body
.as_object_mut()
.expect("a forwarded chat body is an object")
.remove("rid");
assert!(
rid.as_ref()
.and_then(Value::as_str)
.is_some_and(crate::common::is_engine_shaped_rid),
"plain mode must mint an abort rid; got {rid:?}",
);
body
}
async fn assert_forwarded_unchanged(ctx: &Arc<AppContext>, mock: &MockWorker, request: &Value) {
assert_eq!(send(Arc::clone(ctx), request.clone()).await, StatusCode::OK);
assert_eq!(captured(mock), *request);
assert_eq!(without_minted_rid(captured(mock)), *request);
assert!(!ctx
.metrics
.render()
@@ -428,6 +442,6 @@ async fn kimi_ids_forward_with_engine_rendering_fallback() {
} else {
assert!(ids.is_none());
}
assert_eq!(captured(&mock), request);
assert_eq!(without_minted_rid(captured(&mock)), request);
}
}