// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors // SPDX-License-Identifier: Apache-2.0 //! PD pool isolation — end-to-end at the HTTP layer using MockWorker. //! //! Drives the chat handler with: //! //! * A model whose registered workers are all `WorkerMode::Decode`. The //! handler dispatches **prefill** traffic (chat-completions is the //! prefill phase of a PD request), so it must return 503 with //! `no_prefill_workers_available`. //! * A model with no workers at all → 503 `no_healthy_workers` //! (existing code path; pinned here so a future PD wiring change //! doesn't silently swap codes). //! * A PD-disagg model with both pools healthy → request flows to the //! prefill worker (sanity check; the decode worker MUST NOT be selected for //! the chat route). use axum::body::Body; use axum::http::{Request, StatusCode}; use http_body_util::BodyExt; use sgl_router::config::{ Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, }; use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; use sgl_router::policies::factory::build_registry_with_defaults; use sgl_router::proxy::Proxy; use sgl_router::server::app::build_router; use sgl_router::server::app_context::AppContext; use sgl_router::tokenizer::TokenizerRegistry; use sgl_router::workers::WorkerRegistry; use std::sync::Arc; use std::time::Duration; use tower::ServiceExt; fn config() -> Config { Config { server: ServerConfig { host: "0".into(), port: 0, ..Default::default() }, observability: ObservabilityConfig::default(), model: ModelConfig { id: "tiny".into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), disable_input_ids_forwarding: false, policy: PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None, circuit_breaker: None, cache_aware: None, sticky: None, affinity: None, fused: None, eligibility: None, sampling_overrides: Default::default(), }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], }), proxy: ProxyConfig::default(), router_inflight_load: InflightLoadConfig::default(), } } fn build_ctx(specs: Vec) -> Arc { let cfg = config(); let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap()); let registry = Arc::new(WorkerRegistry::default()); for s in specs { let _ = registry.add(s); } let policies = Arc::new(build_registry_with_defaults(&cfg).unwrap()); let proxy = Arc::new(Proxy::new(Duration::from_secs(5)).unwrap()); Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies)) } fn chat_request() -> Request { 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"}], })) .unwrap(), )) .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. #[tokio::test] async fn pd_mode_decode_only_returns_no_prefill_workers_available() { let worker = crate::common::mock_worker::MockWorker::start(vec![]).await; let ctx = build_ctx(vec![WorkerSpec { id: WorkerId("d1".into()), url: worker.url.clone(), mode: WorkerMode::Decode, model_ids: vec![ModelId("tiny".into())], bootstrap_port: None, }]); let app = build_router(ctx); let res = app.oneshot(chat_request()).await.unwrap(); assert_eq!(res.status(), StatusCode::SERVICE_UNAVAILABLE); assert_eq!( res.headers().get("x-router-error-code").unwrap(), "no_prefill_workers_available", ); let body = res.into_body().collect().await.unwrap().to_bytes(); let body_str = String::from_utf8_lossy(&body); assert!( body_str.contains("\"code\":\"no_prefill_workers_available\""), "body: {body_str}" ); } /// Pin the existing-code-path branch: no workers at all → 503 with /// `no_healthy_workers`. Ensures the new PD code path didn't swap the /// code for the "model has zero workers" case. #[tokio::test] async fn no_workers_returns_no_healthy_workers() { let ctx = build_ctx(vec![]); let app = build_router(ctx); let res = app.oneshot(chat_request()).await.unwrap(); assert_eq!(res.status(), StatusCode::SERVICE_UNAVAILABLE); assert_eq!( res.headers().get("x-router-error-code").unwrap(), "no_healthy_workers", ); } /// PD-disagg deployment with both pools healthy → chat dispatch fans /// out to BOTH the prefill and the decode worker (Pattern B: prefill /// in a detached task, decode awaited for the client response). Both /// receive the same bootstrap-injected body so the SGLang engine can /// match KV transfers via `bootstrap_room`. Pool *isolation* — the /// guarantee that the policy's prefill candidate set excludes decode /// workers — is exercised at the resolver layer /// (`policies::registry::tests::pd_resolution_returns_distinct_pools`). /// Here we only assert the HTTP-layer wiring of the dual dispatch. #[tokio::test] async fn pd_mode_chat_dispatch_fans_to_both_prefill_and_decode() { let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await; let decode = crate::common::mock_worker::MockWorker::start(vec![]).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); // Fire a single request; both prefill (spawn-and-forget) and // decode (awaited) must receive a body with the injected // bootstrap fields. The decode body is what the client sees on // the response. let res = app.oneshot(chat_request()).await.unwrap(); assert_eq!( res.status(), StatusCode::OK, "decode response status should reach the client", ); // Decode receives its body synchronously (we awaited it), so it's // guaranteed captured by the time the response returned. Scope // the lock guard to this block so it doesn't span the `.await` // below (clippy: await_holding_lock). { let decode_seen = decode.captured.lock().unwrap(); assert!( decode_seen.last_body.is_some(), "decode worker must receive the bootstrap-injected request body in PD mode", ); } // Prefill is detached; poll briefly until its capture lands. The // prefill task races the HTTP response back to the client. The // local binding releases the `std::sync::Mutex` guard before the // `.await` — holding a sync mutex across an await would let one // task pin the lock while another tries to acquire it. let prefill_body = tokio::time::timeout(Duration::from_secs(2), async { loop { let captured = prefill.captured.lock().unwrap().last_body.clone(); if let Some(b) = captured { return b; } tokio::time::sleep(Duration::from_millis(5)).await; } }) .await .expect("prefill MUST eventually receive its body via the detached task"); assert!(!prefill_body.is_empty()); } /// PD-mode chat request carries an `x-sgl-decode-url` header for the final /// Decode decision. Step 1 defaults to Decode P2; the header remains an /// observability contract regardless of which Decode policy produced it. #[tokio::test] async fn pd_mode_chat_dispatch_sets_final_decode_header() { use std::collections::HashSet; let prefill_a = crate::common::mock_worker::MockWorker::start(vec![]).await; let prefill_b = crate::common::mock_worker::MockWorker::start(vec![]).await; let decode_a = crate::common::mock_worker::MockWorker::start(vec![]).await; let decode_b = crate::common::mock_worker::MockWorker::start(vec![]).await; // MockWorker URLs all bind to `127.0.0.1`; this test deliberately does // not assert a host relation. It pins only the HTTP wiring: the final D // selected by the role-local policy is reflected on the P request. let ctx = build_ctx(vec![ WorkerSpec { id: WorkerId("p1".into()), url: prefill_a.url.clone(), mode: WorkerMode::Prefill, model_ids: vec![ModelId("tiny".into())], bootstrap_port: None, }, WorkerSpec { id: WorkerId("p2".into()), url: prefill_b.url.clone(), mode: WorkerMode::Prefill, model_ids: vec![ModelId("tiny".into())], bootstrap_port: None, }, WorkerSpec { id: WorkerId("d1".into()), url: decode_a.url.clone(), mode: WorkerMode::Decode, model_ids: vec![ModelId("tiny".into())], bootstrap_port: None, }, WorkerSpec { id: WorkerId("d2".into()), url: decode_b.url.clone(), mode: WorkerMode::Decode, model_ids: vec![ModelId("tiny".into())], bootstrap_port: None, }, ]); let app = build_router(ctx); // Fire 4 requests; both prefill workers see traffic via round-robin. for _ in 0..4 { let res = app.clone().oneshot(chat_request()).await.unwrap(); assert_eq!(res.status(), StatusCode::OK); } // Every request that hit a prefill mock MUST carry the final-decode // header. The value MUST be one of the two registered Decode URLs. let decode_urls: HashSet = [decode_a.url.clone(), decode_b.url.clone()] .into_iter() .collect(); for (label, p) in [("prefill_a", &prefill_a), ("prefill_b", &prefill_b)] { let g = p.captured.lock().unwrap(); if g.last_body.is_none() { // This prefill didn't receive a request — round-robin's // dashmap iteration is non-deterministic, so one side may // skip in a 4-request fire. Continue. continue; } let hdr = g.headers.get("x-sgl-decode-url").unwrap_or_else(|| { panic!( "{label} did not receive an x-sgl-decode-url header. headers: {:?}", g.headers ) }); assert!( decode_urls.contains(hdr), "{label} got decode hint {hdr}, expected one of {decode_urls:?}", ); } } /// Task C: plain-mode (non-PD) request does NOT carry the /// `x-sgl-decode-url` header. Pin: the affinity step is gated on /// `worker.mode() == Prefill` so plain workers are not asked to /// bootstrap nonexistent decode peers. #[tokio::test] async fn plain_mode_chat_dispatch_omits_decode_affinity_header() { let plain = crate::common::mock_worker::MockWorker::start(vec![]).await; let ctx = build_ctx(vec![WorkerSpec { id: WorkerId("w1".into()), url: plain.url.clone(), mode: WorkerMode::Plain, model_ids: vec![ModelId("tiny".into())], bootstrap_port: None, }]); let app = build_router(ctx); let res = app.oneshot(chat_request()).await.unwrap(); assert_eq!(res.status(), StatusCode::OK); let g = plain.captured.lock().unwrap(); assert!( !g.headers.contains_key("x-sgl-decode-url"), "plain-mode worker must not receive a decode-affinity header. headers: {:?}", g.headers, ); } /// Task C: PD-mode prefill request with NO decode workers → 503 /// `no_decode_workers_available`. Pin: failure mode is loud and /// distinct from the existing `no_prefill_workers_available` path. #[tokio::test] async fn pd_mode_prefill_only_returns_no_decode_workers_available() { let prefill = crate::common::mock_worker::MockWorker::start(vec![]).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: None, }]); let app = build_router(ctx); let res = app.oneshot(chat_request()).await.unwrap(); assert_eq!(res.status(), StatusCode::SERVICE_UNAVAILABLE); assert_eq!( res.headers().get("x-router-error-code").unwrap(), "no_decode_workers_available", ); } /// PD-mode chat response carries `x-sgl-decode-url` so external tests /// can observe final Decode selection end-to-end (without sniffing the proxy /// hop into the upstream prefill worker). Mirrors the request-side /// behavior asserted by `pd_mode_chat_dispatch_sets_final_decode_header`. #[tokio::test] async fn pd_mode_chat_response_carries_decode_affinity_header() { use std::collections::HashSet; let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await; let decode_a = crate::common::mock_worker::MockWorker::start(vec![]).await; let decode_b = crate::common::mock_worker::MockWorker::start(vec![]).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: None, }, WorkerSpec { id: WorkerId("d1".into()), url: decode_a.url.clone(), mode: WorkerMode::Decode, model_ids: vec![ModelId("tiny".into())], bootstrap_port: None, }, WorkerSpec { id: WorkerId("d2".into()), url: decode_b.url.clone(), mode: WorkerMode::Decode, model_ids: vec![ModelId("tiny".into())], bootstrap_port: None, }, ]); let app = build_router(ctx); let res = app.oneshot(chat_request()).await.unwrap(); assert_eq!(res.status(), StatusCode::OK); let decode_urls: HashSet = [decode_a.url.clone(), decode_b.url.clone()] .into_iter() .collect(); let hdr = res .headers() .get("x-sgl-decode-url") .unwrap_or_else(|| { panic!( "PD-mode chat response did not carry x-sgl-decode-url; headers: {:?}", res.headers(), ) }) .to_str() .unwrap() .to_owned(); assert!( decode_urls.contains(&hdr), "response carried decode hint {hdr}, expected one of {decode_urls:?}", ); } /// Plain-mode chat response does NOT carry `x-sgl-decode-url`. Pin: the /// response-side mirror is gated on PD-mode dispatch. #[tokio::test] async fn plain_mode_chat_response_omits_decode_affinity_header() { let plain = crate::common::mock_worker::MockWorker::start(vec![]).await; let ctx = build_ctx(vec![WorkerSpec { id: WorkerId("w1".into()), url: plain.url.clone(), mode: WorkerMode::Plain, model_ids: vec![ModelId("tiny".into())], bootstrap_port: None, }]); let app = build_router(ctx); let res = app.oneshot(chat_request()).await.unwrap(); assert_eq!(res.status(), StatusCode::OK); assert!( !res.headers().contains_key("x-sgl-decode-url"), "plain-mode chat response must not carry x-sgl-decode-url; headers: {:?}", res.headers(), ); }