// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors // SPDX-License-Identifier: Apache-2.0 //! End-to-end coverage for `policy = "sticky"`: a routing key read from the //! operator-configured header pins a session to one worker, and the //! `sgl_router_sticky_total` outcomes are recorded. Runs against two //! `MockWorker` backends (CPU-only, no GPU). use sgl_router::config::{ Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind, }; use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry; 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 axum::body::Body; use axum::http::{Request, StatusCode}; use std::sync::Arc; use std::time::Duration; use tower::ServiceExt; use crate::common::mock_worker::MockWorker; const TEST_TIMEOUT: Duration = Duration::from_secs(5); /// Build an `AppContext` running the sticky policy over the given worker /// URLs, reading the routing key from `header_name`. Eviction is pushed far /// out so the background sweeper never fires mid-test. fn build_sticky_ctx(header_name: &str, worker_urls: &[String]) -> Arc { let cfg = 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::Sticky, decode_policy: Default::default(), bucket_config: None, circuit_breaker: None, cache_aware: None, sticky: Some(StickyConfig { header_name: header_name.to_string(), fallback_policy: StickyFallbackKind::RoundRobin, idle_secs: 3600, eviction_interval_secs: 3600, }), 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(), }; let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap()); let registry = Arc::new(WorkerRegistry::default()); for (i, url) in worker_urls.iter().enumerate() { let _ = registry.add(WorkerSpec { id: WorkerId(format!("w{i}")), url: url.clone(), mode: WorkerMode::Plain, model_ids: vec![ModelId("tiny".into())], bootstrap_port: None, }); } let policies = Arc::new(build_policy_registry(&cfg).unwrap()); let proxy = Arc::new(Proxy::new(TEST_TIMEOUT).unwrap()); Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies)) } fn chat_request(header: Option<(&str, &str)>) -> Request { let mut builder = Request::builder() .method("POST") .uri("/v1/chat/completions") .header("content-type", "application/json"); if let Some((name, value)) = header { builder = builder.header(name, value); } builder .body(Body::from( serde_json::to_vec(&serde_json::json!({ "model": "tiny", "messages": [{"role": "user", "content": "hi"}], "stream": false })) .unwrap(), )) .unwrap() } /// Parse `sgl_router_worker_requests_total{...,outcome="success"} N` lines into a /// map of worker_url -> success count. fn success_counts(metrics: &str) -> std::collections::HashMap { let mut counts = std::collections::HashMap::new(); for line in metrics.lines() { let Some(rest) = line.strip_prefix("sgl_router_worker_requests_total{") else { continue; }; if !rest.contains(r#"outcome="success""#) { continue; } let Some(url_start) = rest.find(r#"worker_url=""#) else { continue; }; let after = &rest[url_start + r#"worker_url=""#.len()..]; let Some(url_end) = after.find('"') else { continue; }; let url = after[..url_end].to_string(); let value: u64 = line .rsplit(' ') .next() .and_then(|v| v.parse().ok()) .unwrap_or(0); *counts.entry(url).or_insert(0) += value; } counts } /// Read `sgl_router_sticky_total{outcome=""} N`. fn sticky_count(metrics: &str, outcome: &str) -> u64 { let needle = format!(r#"sgl_router_sticky_total{{outcome="{outcome}"}} "#); metrics .lines() .find_map(|l| l.strip_prefix(&needle)) .and_then(|v| v.trim().parse().ok()) .unwrap_or(0) } #[tokio::test] async fn same_routing_key_pins_to_one_worker() { let w0 = MockWorker::start(vec![]).await; let w1 = MockWorker::start(vec![]).await; let ctx = build_sticky_ctx("x-sgl-routing-key", &[w0.url.clone(), w1.url.clone()]); let app = build_router(ctx.clone()); const N: usize = 5; for _ in 0..N { let res = app .clone() .oneshot(chat_request(Some(("x-sgl-routing-key", "alice")))) .await .unwrap(); assert_eq!(res.status(), StatusCode::OK); } let metrics = ctx.metrics.render(); let counts = success_counts(&metrics); let total: u64 = counts.values().sum(); assert_eq!(total, N as u64, "all requests should succeed: {counts:?}"); let pinned: Vec<_> = counts.iter().filter(|(_, &c)| c > 0).collect(); assert_eq!( pinned.len(), 1, "all same-key requests must hit exactly one worker: {counts:?}" ); assert_eq!(*pinned[0].1, N as u64); // One assignment, the rest hits. assert_eq!(sticky_count(&metrics, "assigned"), 1, "metrics:\n{metrics}"); assert_eq!(sticky_count(&metrics, "hit"), (N - 1) as u64); } #[tokio::test] async fn remaps_to_survivor_when_pinned_worker_is_removed() { let w0 = MockWorker::start(vec![]).await; let w1 = MockWorker::start(vec![]).await; let worker_urls = vec![w0.url.clone(), w1.url.clone()]; let ctx = build_sticky_ctx("x-sgl-routing-key", &worker_urls); let app = build_router(ctx.clone()); // Pin the key, then discover which worker it landed on. let res = app .clone() .oneshot(chat_request(Some(("x-sgl-routing-key", "alice")))) .await .unwrap(); assert_eq!(res.status(), StatusCode::OK); let pinned_url = success_counts(&ctx.metrics.render()) .into_iter() .find(|(_, c)| *c > 0) .map(|(url, _)| url) .expect("a worker should have served the first request"); let pinned_idx = worker_urls.iter().position(|u| *u == pinned_url).unwrap(); let survivor_url = worker_urls[1 - pinned_idx].clone(); // Remove the pinned worker from the registry; the next same-key request // must remap to the survivor (not fail). ctx.registry.remove(&WorkerId(format!("w{pinned_idx}"))); let res = app .clone() .oneshot(chat_request(Some(("x-sgl-routing-key", "alice")))) .await .unwrap(); assert_eq!(res.status(), StatusCode::OK); let metrics = ctx.metrics.render(); assert_eq!(sticky_count(&metrics, "remap"), 1, "{metrics}"); // The survivor served the second request. let counts = success_counts(&metrics); assert_eq!(counts.get(&survivor_url).copied().unwrap_or(0), 1); } #[tokio::test] async fn keyless_request_is_served_via_fallback() { let w0 = MockWorker::start(vec![]).await; let w1 = MockWorker::start(vec![]).await; let ctx = build_sticky_ctx("x-sgl-routing-key", &[w0.url.clone(), w1.url.clone()]); let app = build_router(ctx.clone()); // No routing-key header at all. let res = app.clone().oneshot(chat_request(None)).await.unwrap(); assert_eq!(res.status(), StatusCode::OK); let metrics = ctx.metrics.render(); assert_eq!(sticky_count(&metrics, "no_routing_key"), 1, "{metrics}"); assert_eq!(sticky_count(&metrics, "assigned"), 0); } #[tokio::test] async fn only_the_configured_header_name_is_honored() { // Router configured to read the key from `x-session-id`. let w0 = MockWorker::start(vec![]).await; let w1 = MockWorker::start(vec![]).await; let ctx = build_sticky_ctx("x-session-id", &[w0.url.clone(), w1.url.clone()]); let app = build_router(ctx.clone()); // A request using the configured header pins (assigned), and a repeat hits. for _ in 0..2 { let res = app .clone() .oneshot(chat_request(Some(("x-session-id", "s-1")))) .await .unwrap(); assert_eq!(res.status(), StatusCode::OK); } // A request carrying the *default* header (not the configured one) must // be treated as keyless — proving the header name is dynamic, not baked in. let res = app .clone() .oneshot(chat_request(Some(("x-sgl-routing-key", "s-1")))) .await .unwrap(); assert_eq!(res.status(), StatusCode::OK); let metrics = ctx.metrics.render(); assert_eq!(sticky_count(&metrics, "assigned"), 1, "{metrics}"); assert_eq!(sticky_count(&metrics, "hit"), 1); assert_eq!(sticky_count(&metrics, "no_routing_key"), 1); } /// True-sticky scale-up: a worker that joins the registry at runtime must /// NOT redistribute an already-pinned key — the defining difference from /// consistent hashing, where adding a node remaps a fraction of keys. The /// policy unit tests assert this over a bare worker slice; this drives it /// end-to-end through the HTTP stack and a live `WorkerRegistry::add`, so /// the freshly-added worker is a genuine healthy candidate the policy could /// pick — and provably doesn't. #[tokio::test] async fn adding_a_worker_does_not_redistribute_existing_key() { let w0 = MockWorker::start(vec![]).await; let w1 = MockWorker::start(vec![]).await; let ctx = build_sticky_ctx("x-sgl-routing-key", &[w0.url.clone(), w1.url.clone()]); let app = build_router(ctx.clone()); // Pin "alice" to whichever worker the round-robin fallback selects. let res = app .clone() .oneshot(chat_request(Some(("x-sgl-routing-key", "alice")))) .await .unwrap(); assert_eq!(res.status(), StatusCode::OK); let pinned_url = success_counts(&ctx.metrics.render()) .into_iter() .find(|(_, c)| *c > 0) .map(|(url, _)| url) .expect("first request should have been served by some worker"); // Scale up: a third worker joins the registry at runtime. With no // circuit breaker configured it is immediately a healthy candidate // (`healthy_workers_for` filters only on the breaker), so the policy // *could* route to it — the assertions below prove it does not. let w2 = MockWorker::start(vec![]).await; ctx.registry .add(WorkerSpec { id: WorkerId("w2".into()), url: w2.url.clone(), mode: WorkerMode::Plain, model_ids: vec![ModelId("tiny".into())], bootstrap_port: None, }) .unwrap(); // Guard the premise: w2 really is an eligible candidate now, so the // "policy doesn't pick it" assertions below are meaningful and can't // pass vacuously if a future change stops enumerating added workers. assert_eq!( ctx.registry .healthy_workers_for(&ModelId("tiny".into())) .len(), 3, "added worker must be an eligible candidate" ); const N: usize = 5; for _ in 0..N { let res = app .clone() .oneshot(chat_request(Some(("x-sgl-routing-key", "alice")))) .await .unwrap(); assert_eq!(res.status(), StatusCode::OK); } let metrics = ctx.metrics.render(); let counts = success_counts(&metrics); // All N+1 successes stayed on the original pin — nothing leaked to the // newly-added worker or the other pre-existing one. assert_eq!( counts.get(&pinned_url).copied().unwrap_or(0), (N + 1) as u64, "all same-key requests must stay on the original pin: {counts:?}" ); assert_eq!( counts.values().sum::(), (N + 1) as u64, "no request should leak to another worker: {counts:?}" ); // One initial assign, the pin survived the scale-up (no remap), N hits. assert_eq!(sticky_count(&metrics, "assigned"), 1, "{metrics}"); assert_eq!(sticky_count(&metrics, "remap"), 0, "{metrics}"); assert_eq!(sticky_count(&metrics, "hit"), N as u64, "{metrics}"); }