[model-gateway] : Rust integration tests for integration_mock replacement (#16441)

sgl-model-gateway/tests/routing/cache_aware_backward_compat_test.rs

@@ -0,0 +1,165 @@
+use std::{collections::HashMap, sync::Arc};
+
+use smg::{
+    core::{BasicWorkerBuilder, Worker, WorkerType},
+    policies::{CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, SelectWorkerInfo},
+};
+
+#[test]
+fn test_backward_compatibility_with_empty_model_id() {
+    let config = CacheAwareConfig {
+        cache_threshold: 0.5,
+        balance_abs_threshold: 2,
+        balance_rel_threshold: 1.5,
+        eviction_interval_secs: 0, // Disable background eviction for testing
+        max_tree_size: 100,
+    };
+
+    let policy = CacheAwarePolicy::with_config(config);
+
+    // Create workers with empty model_id (simulating existing routers)
+    let worker1 = BasicWorkerBuilder::new("http://worker1:8080")
+        .worker_type(WorkerType::Regular)
+        .api_key("test_api_key")
+        .build();
+    // No model_id label - should default to "unknown"
+
+    let mut labels2 = HashMap::new();
+    labels2.insert("model_id".to_string(), "unknown".to_string());
+    let worker2 = BasicWorkerBuilder::new("http://worker2:8080")
+        .worker_type(WorkerType::Regular)
+        .api_key("test_api_key")
+        .labels(labels2)
+        .build();
+
+    // Add workers - should both go to "default" tree
+    policy.add_worker(&worker1);
+    policy.add_worker(&worker2);
+
+    // Create worker list
+    let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(worker1.clone()), Arc::new(worker2.clone())];
+
+    // Select worker - should work without errors
+    let selected = policy.select_worker(
+        &workers,
+        &SelectWorkerInfo {
+            request_text: Some("test request"),
+            ..Default::default()
+        },
+    );
+    assert!(selected.is_some(), "Should select a worker");
+
+    // Remove workers - should work without errors
+    policy.remove_worker(&worker1);
+    policy.remove_worker(&worker2);
+}
+
+#[test]
+fn test_mixed_model_ids() {
+    let config = CacheAwareConfig {
+        cache_threshold: 0.5,
+        balance_abs_threshold: 2,
+        balance_rel_threshold: 1.5,
+        eviction_interval_secs: 0,
+        max_tree_size: 100,
+    };
+
+    let policy = CacheAwarePolicy::with_config(config);
+
+    // Create workers with different model_id scenarios
+    let worker1 = BasicWorkerBuilder::new("http://worker1:8080")
+        .worker_type(WorkerType::Regular)
+        .api_key("test_api_key")
+        .build();
+    // No model_id label - defaults to "unknown" which goes to "default" tree
+
+    let mut labels2 = HashMap::new();
+    labels2.insert("model_id".to_string(), "llama-3".to_string());
+    let worker2 = BasicWorkerBuilder::new("http://worker2:8080")
+        .worker_type(WorkerType::Regular)
+        .labels(labels2)
+        .api_key("test_api_key")
+        .build();
+
+    let mut labels3 = HashMap::new();
+    labels3.insert("model_id".to_string(), "unknown".to_string());
+    let worker3 = BasicWorkerBuilder::new("http://worker3:8080")
+        .worker_type(WorkerType::Regular)
+        .labels(labels3)
+        .build();
+
+    let mut labels4 = HashMap::new();
+    labels4.insert("model_id".to_string(), "llama-3".to_string());
+    let worker4 = BasicWorkerBuilder::new("http://worker4:8080")
+        .worker_type(WorkerType::Regular)
+        .labels(labels4)
+        .build();
+
+    // Add all workers
+    policy.add_worker(&worker1);
+    policy.add_worker(&worker2);
+    policy.add_worker(&worker3);
+    policy.add_worker(&worker4);
+
+    let default_workers: Vec<Arc<dyn Worker>> =
+        vec![Arc::new(worker1.clone()), Arc::new(worker3.clone())];
+    let info = SelectWorkerInfo {
+        request_text: Some("test request"),
+        ..Default::default()
+    };
+    let selected = policy.select_worker(&default_workers, &info);
+    assert!(selected.is_some(), "Should select from default workers");
+
+    let llama_workers: Vec<Arc<dyn Worker>> =
+        vec![Arc::new(worker2.clone()), Arc::new(worker4.clone())];
+    let selected = policy.select_worker(&llama_workers, &info);
+    assert!(selected.is_some(), "Should select from llama-3 workers");
+
+    let all_workers: Vec<Arc<dyn Worker>> = vec![
+        Arc::new(worker1.clone()),
+        Arc::new(worker2.clone()),
+        Arc::new(worker3.clone()),
+        Arc::new(worker4.clone()),
+    ];
+    let selected = policy.select_worker(&all_workers, &info);
+    assert!(selected.is_some(), "Should select from all workers");
+}
+
+#[test]
+fn test_remove_worker_by_url_backward_compat() {
+    let config = CacheAwareConfig::default();
+    let policy = CacheAwarePolicy::with_config(config);
+
+    // Create workers with different model_ids
+    let mut labels1 = HashMap::new();
+    labels1.insert("model_id".to_string(), "llama-3".to_string());
+    let worker1 = BasicWorkerBuilder::new("http://worker1:8080")
+        .worker_type(WorkerType::Regular)
+        .labels(labels1)
+        .api_key("test_api_key")
+        .build();
+
+    let worker2 = BasicWorkerBuilder::new("http://worker2:8080")
+        .worker_type(WorkerType::Regular)
+        .api_key("test_api_key")
+        .build();
+    // No model_id label - defaults to "unknown"
+
+    // Add workers
+    policy.add_worker(&worker1);
+    policy.add_worker(&worker2);
+
+    // Remove by URL (backward compatibility method)
+    // Should remove from all trees since we don't know the model
+    policy.remove_worker_by_url("http://worker1:8080");
+
+    let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(worker2.clone())];
+    let selected = policy.select_worker(
+        &workers,
+        &SelectWorkerInfo {
+            request_text: Some("test"),
+            ..Default::default()
+        },
+    );
+    assert_eq!(selected, Some(0), "Should only have worker2 left");
+}

sgl-model-gateway/tests/routing/load_balancing_test.rs

@@ -0,0 +1,373 @@
+//! Load balancing integration tests
+//!
+//! Tests for various load balancing policies: round_robin, random, cache_aware, etc.
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use smg::config::{CircuitBreakerConfig, RetryConfig};
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+#[cfg(test)]
+mod round_robin_tests {
+    use super::*;
+
+    /// Test that round robin distributes requests evenly across workers
+    #[tokio::test]
+    async fn test_round_robin_distribution() {
+        let config = TestRouterConfig::round_robin(3100);
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19001, 3))
+                .await;
+
+        let app = ctx.create_app().await;
+        let num_requests = 30;
+        let mut success_count = 0;
+
+        for i in 0..num_requests {
+            let payload = json!({
+                "text": format!("Request {}", i),
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // All requests should succeed with round robin across 3 healthy workers
+        assert_eq!(
+            success_count, num_requests,
+            "All requests should succeed with round robin"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test round robin with one worker failing
+    #[tokio::test]
+    async fn test_round_robin_with_failing_worker() {
+        let config = TestRouterConfig::round_robin_with_retry(
+            3101,
+            RetryConfig {
+                max_retries: 2,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 50,
+                ..Default::default()
+            },
+        );
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(19004, 1.0), // Always fail
+                TestWorkerConfig::healthy(19005),    // Always succeed
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // With retries enabled, requests should eventually succeed
+        // by being retried on the healthy worker
+        let payload = json!({
+            "text": "Test with retries",
+            "stream": false
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/generate")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Request should succeed via retry on healthy worker"
+        );
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod random_tests {
+    use super::*;
+
+    /// Test that random policy distributes requests across workers
+    #[tokio::test]
+    async fn test_random_distribution() {
+        let config = TestRouterConfig::random(3102);
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19010, 2))
+                .await;
+
+        let app = ctx.create_app().await;
+        let num_requests = 20;
+        let mut success_count = 0;
+
+        for i in 0..num_requests {
+            let payload = json!({
+                "text": format!("Random request {}", i),
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // All requests should succeed with random policy
+        assert_eq!(
+            success_count, num_requests,
+            "All requests should succeed with random policy"
+        );
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod cache_aware_tests {
+    use super::*;
+
+    /// Test cache-aware routing uses consistent hashing
+    #[tokio::test]
+    async fn test_cache_aware_consistent_routing() {
+        let config = TestRouterConfig::cache_aware(3103);
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19020, 2))
+                .await;
+
+        let app = ctx.create_app().await;
+
+        // Same prompt should route to same worker (consistent hashing)
+        let same_prompt = "Hello, cache-aware routing test!";
+        let mut worker_ids: Vec<Option<String>> = Vec::new();
+
+        for _ in 0..5 {
+            let payload = json!({
+                "text": same_prompt,
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            assert_eq!(resp.status(), StatusCode::OK);
+
+            // Extract worker ID from response header if available
+            let worker_id = resp
+                .headers()
+                .get("x-worker-id")
+                .map(|v| v.to_str().unwrap().to_string());
+            worker_ids.push(worker_id);
+        }
+
+        // All requests should succeed
+        assert_eq!(worker_ids.len(), 5);
+
+        ctx.shutdown().await;
+    }
+
+    /// Test cache-aware routing with different prompts
+    #[tokio::test]
+    async fn test_cache_aware_different_prompts() {
+        let config = TestRouterConfig::cache_aware(3104);
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19022, 2))
+                .await;
+
+        let app = ctx.create_app().await;
+        let prompts = vec![
+            "First unique prompt",
+            "Second unique prompt",
+            "Third unique prompt",
+            "Fourth unique prompt",
+        ];
+
+        let mut success_count = 0;
+        for prompt in prompts {
+            let payload = json!({
+                "text": prompt,
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // All requests should succeed
+        assert_eq!(success_count, 4);
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod worker_health_tests {
+    use super::*;
+
+    /// Test that failing workers are avoided via circuit breaker/retry
+    #[tokio::test]
+    async fn test_skip_failing_workers() {
+        let config = TestRouterConfig::round_robin_with_reliability(
+            3105,
+            RetryConfig {
+                max_retries: 3,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 50,
+                ..Default::default()
+            },
+            CircuitBreakerConfig {
+                failure_threshold: 2,
+                success_threshold: 1,
+                timeout_duration_secs: 2,
+                window_duration_secs: 10,
+            },
+        );
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(19030, 1.0), // Always fails
+                TestWorkerConfig::healthy(19031),    // Always succeeds
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Requests should succeed by retrying on healthy worker
+        // or by circuit breaker opening on failing worker
+        for i in 0..10 {
+            let payload = json!({
+                "text": format!("Request {}", i),
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            assert_eq!(resp.status(), StatusCode::OK);
+        }
+
+        ctx.shutdown().await;
+    }
+
+    /// Test behavior when all workers are unhealthy
+    #[tokio::test]
+    async fn test_all_workers_unhealthy() {
+        let ctx = AppTestContext::new(vec![]).await;
+
+        let app = ctx.create_app().await;
+
+        let payload = json!({
+            "text": "Test with no workers",
+            "stream": false
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/generate")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        // Should return service unavailable when no workers available
+        assert!(
+            resp.status() == StatusCode::SERVICE_UNAVAILABLE
+                || resp.status() == StatusCode::INTERNAL_SERVER_ERROR,
+            "Expected 503 or 500 when no workers available, got {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod worker_response_delay_tests {
+    use super::*;
+
+    /// Test that response delays are handled correctly
+    #[tokio::test]
+    async fn test_worker_with_delay() {
+        let config = TestRouterConfig::random(3106);
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![TestWorkerConfig::slow(19040, 100)], // 100ms delay
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        let start = std::time::Instant::now();
+        let payload = json!({
+            "text": "Test with delay",
+            "stream": false
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/generate")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        let elapsed = start.elapsed();
+
+        assert_eq!(resp.status(), StatusCode::OK);
+        // Response should take at least 100ms due to configured delay
+        assert!(
+            elapsed.as_millis() >= 100,
+            "Response should be delayed by at least 100ms, got {}ms",
+            elapsed.as_millis()
+        );
+
+        ctx.shutdown().await;
+    }
+}

sgl-model-gateway/tests/routing/manual_routing_test.rs

@@ -0,0 +1,174 @@
+//! Manual routing policy integration tests
+//!
+//! Tests for the manual routing policy with sticky sessions using X-SMG-Routing-Key header.
+
+use std::collections::{HashMap, HashSet};
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+const ROUTING_KEY_HEADER: &str = "X-SMG-Routing-Key";
+
+#[cfg(test)]
+mod manual_routing_tests {
+    use super::*;
+
+    /// Test sticky routing with X-SMG-Routing-Key header
+    #[tokio::test]
+    async fn test_manual_routing_with_header() {
+        let config = TestRouterConfig::manual(3700);
+
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19700, 2))
+                .await;
+
+        let app = ctx.create_app().await;
+
+        // Send requests with different routing keys
+        let mut key_workers: HashMap<String, HashSet<String>> = HashMap::new();
+
+        for key_id in 0..5 {
+            let routing_key = format!("user-{}", key_id);
+
+            for _ in 0..4 {
+                let payload = json!({
+                    "text": format!("Request for {}", routing_key),
+                    "stream": false
+                });
+
+                let req = Request::builder()
+                    .method("POST")
+                    .uri("/generate")
+                    .header(CONTENT_TYPE, "application/json")
+                    .header(ROUTING_KEY_HEADER, &routing_key)
+                    .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                    .unwrap();
+
+                let resp = app.clone().oneshot(req).await.unwrap();
+                assert_eq!(resp.status(), StatusCode::OK);
+
+                // Extract worker ID from response header
+                if let Some(worker_id) = resp.headers().get("x-worker-id") {
+                    let worker = worker_id.to_str().unwrap().to_string();
+                    key_workers
+                        .entry(routing_key.clone())
+                        .or_default()
+                        .insert(worker);
+                }
+            }
+        }
+
+        // Verify sticky: each routing key should route to exactly one worker
+        for (key, workers) in &key_workers {
+            assert_eq!(
+                workers.len(),
+                1,
+                "Routing key {} should route to exactly one worker, got {:?}",
+                key,
+                workers
+            );
+        }
+
+        ctx.shutdown().await;
+    }
+
+    /// Test random fallback when no routing key header is provided
+    #[tokio::test]
+    async fn test_manual_routing_without_header() {
+        let config = TestRouterConfig::manual(3701);
+
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19702, 2))
+                .await;
+
+        let app = ctx.create_app().await;
+        let mut success_count = 0;
+
+        // Send requests without routing key - should fall back to random selection
+        for i in 0..20 {
+            let payload = json!({
+                "text": format!("Request without key {}", i),
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                // No ROUTING_KEY_HEADER
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // All requests should succeed via random fallback
+        assert_eq!(
+            success_count, 20,
+            "All requests should succeed with random fallback when no routing key provided"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test that same routing key consistently routes to same worker
+    #[tokio::test]
+    async fn test_manual_routing_consistency() {
+        let config = TestRouterConfig::manual(3702);
+
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19704, 3))
+                .await;
+
+        let app = ctx.create_app().await;
+
+        let routing_key = "consistent-user-123";
+        let mut seen_workers: Vec<String> = Vec::new();
+
+        // Send multiple requests with same routing key
+        for i in 0..10 {
+            let payload = json!({
+                "text": format!("Consistent request {}", i),
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .header(ROUTING_KEY_HEADER, routing_key)
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            assert_eq!(resp.status(), StatusCode::OK);
+
+            if let Some(worker_id) = resp.headers().get("x-worker-id") {
+                seen_workers.push(worker_id.to_str().unwrap().to_string());
+            }
+        }
+
+        // All requests should go to the same worker
+        if !seen_workers.is_empty() {
+            let first_worker = &seen_workers[0];
+            for worker in &seen_workers {
+                assert_eq!(
+                    worker, first_worker,
+                    "All requests with same routing key should go to same worker"
+                );
+            }
+        }
+
+        ctx.shutdown().await;
+    }
+}

sgl-model-gateway/tests/routing/mod.rs

@@ -0,0 +1,9 @@
+//! Routing integration tests
+
+pub mod cache_aware_backward_compat_test;
+pub mod load_balancing_test;
+pub mod manual_routing_test;
+pub mod pd_routing_test;
+pub mod power_of_two_test;
+pub mod test_openai_routing;
+pub mod test_pd_routing;

sgl-model-gateway/tests/routing/pd_routing_test.rs

@@ -0,0 +1,222 @@
+//! Prefill/Decode (PD) routing integration tests
+//!
+//! Tests for prefill-decode disaggregation routing mode.
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use smg::config::RouterConfig;
+use tower::ServiceExt;
+
+use crate::common::{
+    mock_worker::{HealthStatus, MockWorkerConfig, WorkerType},
+    AppTestContext, TestWorkerConfig,
+};
+
+#[cfg(test)]
+mod pd_routing_tests {
+    use super::*;
+
+    /// Test basic PD mode routing with prefill and decode workers
+    #[tokio::test]
+    async fn test_pd_mode_basic_routing() {
+        let config = RouterConfig::builder()
+            .prefill_decode_mode(
+                vec![
+                    ("http://127.0.0.1:19800".to_string(), None),
+                    ("http://127.0.0.1:19801".to_string(), None),
+                ],
+                vec![
+                    "http://127.0.0.1:19802".to_string(),
+                    "http://127.0.0.1:19803".to_string(),
+                ],
+            )
+            .power_of_two_policy(1)
+            .host("127.0.0.1")
+            .port(3800)
+            .max_payload_size(256 * 1024 * 1024)
+            .request_timeout_secs(600)
+            .worker_startup_timeout_secs(5)
+            .worker_startup_check_interval_secs(1)
+            .max_concurrent_requests(64)
+            .queue_timeout_secs(60)
+            .build_unchecked();
+
+        // Note: For PD mode tests, we need to start prefill and decode workers separately
+        // The test context will need to handle this specially
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                // Prefill workers
+                TestWorkerConfig::prefill(19800),
+                TestWorkerConfig::prefill(19801),
+                // Decode workers
+                TestWorkerConfig::decode(19802),
+                TestWorkerConfig::decode(19803),
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Send requests and verify they succeed
+        for i in 0..10 {
+            let payload = json!({
+                "text": format!("PD mode request {}", i),
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            assert_eq!(
+                resp.status(),
+                StatusCode::OK,
+                "PD mode request should succeed"
+            );
+        }
+
+        ctx.shutdown().await;
+    }
+
+    /// Test PD mode with round robin policy
+    #[tokio::test]
+    async fn test_pd_mode_round_robin() {
+        let config = RouterConfig::builder()
+            .prefill_decode_mode(
+                vec![("http://127.0.0.1:19810".to_string(), None)],
+                vec![
+                    "http://127.0.0.1:19811".to_string(),
+                    "http://127.0.0.1:19812".to_string(),
+                ],
+            )
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(3801)
+            .max_payload_size(256 * 1024 * 1024)
+            .request_timeout_secs(600)
+            .worker_startup_timeout_secs(5)
+            .worker_startup_check_interval_secs(1)
+            .max_concurrent_requests(64)
+            .queue_timeout_secs(60)
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::prefill(19810),
+                TestWorkerConfig::decode(19811),
+                TestWorkerConfig::decode(19812),
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+        let mut success_count = 0;
+
+        for i in 0..20 {
+            let payload = json!({
+                "text": format!("PD round robin {}", i),
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        assert_eq!(
+            success_count, 20,
+            "All requests should succeed in PD mode with round robin"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test PD mode handles worker failures gracefully
+    #[tokio::test]
+    async fn test_pd_mode_with_failing_decode_worker() {
+        use smg::config::RetryConfig;
+
+        let config = RouterConfig::builder()
+            .prefill_decode_mode(
+                vec![("http://127.0.0.1:19820".to_string(), None)],
+                vec![
+                    "http://127.0.0.1:19821".to_string(),
+                    "http://127.0.0.1:19822".to_string(),
+                ],
+            )
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(3802)
+            .max_payload_size(256 * 1024 * 1024)
+            .request_timeout_secs(600)
+            .worker_startup_timeout_secs(5)
+            .worker_startup_check_interval_secs(1)
+            .max_concurrent_requests(64)
+            .queue_timeout_secs(60)
+            .retry_config(RetryConfig {
+                max_retries: 3,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 50,
+                ..Default::default()
+            })
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::prefill(19820),
+                MockWorkerConfig {
+                    port: 19821,
+                    worker_type: WorkerType::Decode,
+                    health_status: HealthStatus::Healthy,
+                    response_delay_ms: 0,
+                    fail_rate: 1.0, // Failing decode worker
+                },
+                TestWorkerConfig::decode(19822), // Healthy decode worker
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Request should succeed via retry to healthy decode worker
+        let payload = json!({
+            "text": "Test with failing decode worker",
+            "stream": false
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/generate")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Request should succeed via retry to healthy decode worker"
+        );
+
+        ctx.shutdown().await;
+    }
+}

sgl-model-gateway/tests/routing/power_of_two_test.rs

@@ -0,0 +1,193 @@
+//! Power of Two load balancing integration tests
+//!
+//! Tests for the Power of Two Choices algorithm that selects the less loaded worker.
+
+use std::sync::{
+    atomic::{AtomicUsize, Ordering},
+    Arc,
+};
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use smg::config::RouterConfig;
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+#[cfg(test)]
+mod power_of_two_tests {
+    use super::*;
+
+    /// Test that power of two distributes requests across workers
+    #[tokio::test]
+    async fn test_power_of_two_distribution() {
+        let config = TestRouterConfig::power_of_two(3600);
+
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19600, 2))
+                .await;
+
+        let app = ctx.create_app().await;
+        let num_requests = 20;
+        let mut success_count = 0;
+
+        for i in 0..num_requests {
+            let payload = json!({
+                "text": format!("Power of two request {}", i),
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // All requests should succeed
+        assert_eq!(
+            success_count, num_requests,
+            "All requests should succeed with power of two policy"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test that power of two prefers less loaded workers
+    #[tokio::test]
+    async fn test_power_of_two_prefers_less_loaded() {
+        let config = TestRouterConfig::power_of_two(3601);
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::slow(19602, 200), // Slow worker
+                TestWorkerConfig::healthy(19603),   // Fast worker
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Send concurrent requests to create load imbalance
+        let mut handles = Vec::new();
+        let success_count = Arc::new(AtomicUsize::new(0));
+
+        for i in 0..30 {
+            let app_clone = app.clone();
+            let success_clone = Arc::clone(&success_count);
+
+            let handle = tokio::spawn(async move {
+                let payload = json!({
+                    "text": format!("Concurrent request {}", i),
+                    "stream": false
+                });
+
+                let req = Request::builder()
+                    .method("POST")
+                    .uri("/generate")
+                    .header(CONTENT_TYPE, "application/json")
+                    .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                    .unwrap();
+
+                let resp = app_clone.oneshot(req).await.unwrap();
+                if resp.status() == StatusCode::OK {
+                    success_clone.fetch_add(1, Ordering::SeqCst);
+                }
+            });
+
+            handles.push(handle);
+        }
+
+        for handle in handles {
+            handle.await.unwrap();
+        }
+
+        // All requests should succeed
+        assert!(
+            success_count.load(Ordering::SeqCst) >= 25,
+            "Most requests should succeed with power of two"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test power of two with failing worker uses retry/CB to route to healthy worker
+    #[tokio::test]
+    async fn test_power_of_two_with_failing_worker() {
+        use smg::config::{CircuitBreakerConfig, RetryConfig};
+
+        let retry_config = RetryConfig {
+            max_retries: 3,
+            initial_backoff_ms: 10,
+            max_backoff_ms: 50,
+            ..Default::default()
+        };
+        let circuit_breaker = CircuitBreakerConfig {
+            failure_threshold: 2,
+            success_threshold: 1,
+            timeout_duration_secs: 2,
+            window_duration_secs: 10,
+        };
+
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .power_of_two_policy(1)
+            .host("127.0.0.1")
+            .port(3602)
+            .max_payload_size(256 * 1024 * 1024)
+            .request_timeout_secs(600)
+            .worker_startup_timeout_secs(5)
+            .worker_startup_check_interval_secs(1)
+            .max_concurrent_requests(64)
+            .queue_timeout_secs(60)
+            .retry_config(retry_config)
+            .circuit_breaker_config(circuit_breaker)
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(19604, 1.0), // Always fails
+                TestWorkerConfig::healthy(19605),    // Always succeeds
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Requests should succeed via retry to healthy worker
+        for i in 0..10 {
+            let payload = json!({
+                "text": format!("Request with failing worker {}", i),
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            assert_eq!(
+                resp.status(),
+                StatusCode::OK,
+                "Request should succeed via retry on healthy worker"
+            );
+        }
+
+        ctx.shutdown().await;
+    }
+}

sgl-model-gateway/tests/routing/test_openai_routing.rs

@@ -0,0 +1,905 @@
+//! Comprehensive integration tests for OpenAI backend functionality
+
+use std::{
+    collections::HashMap,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+};
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{Method, StatusCode},
+    response::Response,
+    routing::post,
+    Json, Router,
+};
+use serde_json::json;
+use smg::{
+    config::{
+        ConfigError, ConfigValidator, HistoryBackend, OracleConfig, RouterConfig, RoutingMode,
+    },
+    data_connector::{ResponseId, StoredResponse},
+    protocols::{
+        chat::{ChatCompletionRequest, ChatMessage, MessageContent},
+        common::StringOrArray,
+        completion::CompletionRequest,
+        generate::GenerateRequest,
+        responses::{ResponseInput, ResponsesGetParams, ResponsesRequest},
+    },
+    routers::{openai::OpenAIRouter, RouterTrait},
+};
+use tokio::{
+    net::TcpListener,
+    time::{sleep, Duration},
+};
+use tower::ServiceExt;
+
+use crate::common::mock_openai_server::MockOpenAIServer;
+
+/// Helper function to create a minimal chat completion request for testing
+fn create_minimal_chat_request() -> ChatCompletionRequest {
+    let val = json!({
+        "model": "gpt-3.5-turbo",
+        "messages": [
+            {"role": "user", "content": "Hello"}
+        ],
+        "max_tokens": 100
+    });
+    serde_json::from_value(val).unwrap()
+}
+
+/// Helper function to create a minimal completion request for testing
+fn create_minimal_completion_request() -> CompletionRequest {
+    CompletionRequest {
+        model: "gpt-3.5-turbo".to_string(),
+        prompt: StringOrArray::String("Hello".to_string()),
+        suffix: None,
+        max_tokens: Some(100),
+        temperature: None,
+        top_p: None,
+        n: None,
+        stream: false,
+        stream_options: None,
+        logprobs: None,
+        echo: false,
+        stop: None,
+        presence_penalty: None,
+        frequency_penalty: None,
+        best_of: None,
+        logit_bias: None,
+        user: None,
+        seed: None,
+        top_k: None,
+        min_p: None,
+        min_tokens: None,
+        repetition_penalty: None,
+        regex: None,
+        ebnf: None,
+        json_schema: None,
+        stop_token_ids: None,
+        no_stop_trim: false,
+        ignore_eos: false,
+        skip_special_tokens: true,
+        lora_path: None,
+        session_params: None,
+        return_hidden_states: false,
+        sampling_seed: None,
+        other: serde_json::Map::new(),
+    }
+}
+
+/// Test basic OpenAI router creation and configuration
+#[tokio::test]
+async fn test_openai_router_creation() {
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    // Register an external worker before creating the router
+    crate::common::test_app::register_external_worker(&ctx, "https://api.openai.com", None);
+    let router = OpenAIRouter::new(&ctx).await;
+
+    assert!(router.is_ok(), "Router creation should succeed");
+
+    let router = router.unwrap();
+    assert_eq!(router.router_type(), "openai");
+    assert!(!router.is_pd_mode());
+}
+
+/// Test server info endpoint
+#[tokio::test]
+async fn test_openai_router_server_info() {
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    crate::common::test_app::register_external_worker(&ctx, "https://api.openai.com", None);
+    let router = OpenAIRouter::new(&ctx).await.unwrap();
+
+    let req = Request::builder()
+        .method(Method::GET)
+        .uri("/info")
+        .body(Body::empty())
+        .unwrap();
+
+    let response = router.get_server_info(req).await;
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let (_, body) = response.into_parts();
+    let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap();
+    let body_str = String::from_utf8(body_bytes.to_vec()).unwrap();
+
+    assert!(body_str.contains("openai"));
+}
+
+/// Test models endpoint
+#[tokio::test]
+async fn test_openai_router_models() {
+    // Use mock server for deterministic models response
+    let mock_server = MockOpenAIServer::new().await;
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    crate::common::test_app::register_external_worker(&ctx, &mock_server.base_url(), None);
+    let router = OpenAIRouter::new(&ctx).await.unwrap();
+
+    let req = Request::builder()
+        .method(Method::GET)
+        .uri("/models")
+        .body(Body::empty())
+        .unwrap();
+
+    let response = router.get_models(req).await;
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let (_, body) = response.into_parts();
+    let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap();
+    let body_str = String::from_utf8(body_bytes.to_vec()).unwrap();
+    let models: serde_json::Value = serde_json::from_str(&body_str).unwrap();
+
+    assert_eq!(models["object"], "list");
+    assert!(models["data"].is_array());
+}
+
+#[tokio::test]
+async fn test_openai_router_responses_with_mock() {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    let counter = Arc::new(AtomicUsize::new(0));
+    let counter_clone = counter.clone();
+
+    let app = Router::new().route(
+        "/v1/responses",
+        post({
+            move |Json(request): Json<serde_json::Value>| {
+                let counter = counter_clone.clone();
+                async move {
+                    let idx = counter.fetch_add(1, Ordering::SeqCst) + 1;
+                    let model = request
+                        .get("model")
+                        .and_then(|v| v.as_str())
+                        .unwrap_or("gpt-4o-mini")
+                        .to_string();
+                    let id = format!("resp_mock_{idx}");
+                    let response = json!({
+                        "id": id,
+                        "object": "response",
+                        "created_at": 1_700_000_000 + idx as i64,
+                        "status": "completed",
+                        "model": model,
+                        "output": [{
+                            "type": "message",
+                            "id": format!("msg_{idx}"),
+                            "role": "assistant",
+                            "status": "completed",
+                            "content": [{
+                                "type": "output_text",
+                                "text": format!("mock_output_{idx}"),
+                                "annotations": []
+                            }]
+                        }],
+                        "metadata": {}
+                    });
+                    Json(response)
+                }
+            }
+        }),
+    );
+
+    let server = tokio::spawn(async move {
+        axum::serve(listener, app).await.unwrap();
+    });
+
+    let base_url = format!("http://{}", addr);
+
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    crate::common::test_app::register_external_worker(&ctx, &base_url, Some(vec!["gpt-4o-mini"]));
+    let router = OpenAIRouter::new(&ctx).await.unwrap();
+
+    // Get storage from context (router uses this, not a separate storage)
+    let storage = ctx.response_storage.clone();
+
+    let request1 = ResponsesRequest {
+        model: "gpt-4o-mini".to_string(),
+        input: ResponseInput::Text("Say hi".to_string()),
+        store: Some(true),
+        ..Default::default()
+    };
+
+    let response1 = router.route_responses(None, &request1, None).await;
+    assert_eq!(response1.status(), StatusCode::OK);
+    let body1_bytes = axum::body::to_bytes(response1.into_body(), usize::MAX)
+        .await
+        .unwrap();
+    let body1: serde_json::Value = serde_json::from_slice(&body1_bytes).unwrap();
+    let resp1_id = body1["id"].as_str().expect("id missing").to_string();
+    assert_eq!(body1["previous_response_id"], serde_json::Value::Null);
+
+    let request2 = ResponsesRequest {
+        model: "gpt-4o-mini".to_string(),
+        input: ResponseInput::Text("Thanks".to_string()),
+        store: Some(true),
+        previous_response_id: Some(resp1_id.clone()),
+        ..Default::default()
+    };
+
+    let response2 = router.route_responses(None, &request2, None).await;
+    assert_eq!(response2.status(), StatusCode::OK);
+    let body2_bytes = axum::body::to_bytes(response2.into_body(), usize::MAX)
+        .await
+        .unwrap();
+    let body2: serde_json::Value = serde_json::from_slice(&body2_bytes).unwrap();
+    let resp2_id = body2["id"].as_str().expect("second id missing");
+    assert_eq!(
+        body2["previous_response_id"].as_str(),
+        Some(resp1_id.as_str())
+    );
+
+    let stored1 = storage
+        .get_response(&ResponseId::from(resp1_id.clone()))
+        .await
+        .unwrap()
+        .expect("first response missing");
+    // Input is now stored as a JSON array of items
+    assert!(stored1.input.is_array());
+    let input_items = stored1.input.as_array().unwrap();
+    assert_eq!(input_items.len(), 1);
+    assert_eq!(input_items[0]["type"], "message");
+    assert_eq!(input_items[0]["role"], "user");
+    assert_eq!(input_items[0]["content"][0]["text"], "Say hi");
+
+    // Output is now stored as a JSON array of items
+    assert!(stored1.output.is_array());
+    let output_items = stored1.output.as_array().unwrap();
+    assert_eq!(output_items.len(), 1);
+    assert_eq!(output_items[0]["content"][0]["text"], "mock_output_1");
+
+    assert!(stored1.previous_response_id.is_none());
+
+    let stored2 = storage
+        .get_response(&ResponseId::from(resp2_id))
+        .await
+        .unwrap()
+        .expect("second response missing");
+    assert_eq!(stored2.previous_response_id.unwrap().0, resp1_id);
+
+    // Output is now stored as a JSON array
+    assert!(stored2.output.is_array());
+    let output_items2 = stored2.output.as_array().unwrap();
+    assert_eq!(output_items2.len(), 1);
+    assert_eq!(output_items2[0]["content"][0]["text"], "mock_output_2");
+
+    let get1 = router
+        .get_response(None, &stored1.id.0, &ResponsesGetParams::default())
+        .await;
+    assert_eq!(get1.status(), StatusCode::OK);
+    let get1_body_bytes = axum::body::to_bytes(get1.into_body(), usize::MAX)
+        .await
+        .unwrap();
+    let get1_json: serde_json::Value = serde_json::from_slice(&get1_body_bytes).unwrap();
+    assert_eq!(get1_json, body1);
+
+    let get2 = router
+        .get_response(None, &stored2.id.0, &ResponsesGetParams::default())
+        .await;
+    assert_eq!(get2.status(), StatusCode::OK);
+    let get2_body_bytes = axum::body::to_bytes(get2.into_body(), usize::MAX)
+        .await
+        .unwrap();
+    let get2_json: serde_json::Value = serde_json::from_slice(&get2_body_bytes).unwrap();
+    assert_eq!(get2_json, body2);
+
+    server.abort();
+}
+
+#[tokio::test]
+async fn test_openai_router_responses_streaming_with_mock() {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+
+    let sse_handler = post(|Json(_request): Json<serde_json::Value>| async move {
+        let response_id = "resp_stream_123";
+        let message_id = "msg_stream_123";
+        let final_text = "Once upon a streamed unicorn adventure.";
+
+        let events = vec![
+            (
+                "response.created",
+                json!({
+                    "type": "response.created",
+                    "sequence_number": 0,
+                    "response": {
+                        "id": response_id,
+                        "object": "response",
+                        "created_at": 1_700_000_500,
+                        "status": "in_progress",
+                        "model": "",
+                        "output": [],
+                        "parallel_tool_calls": true,
+                        "previous_response_id": null,
+                        "reasoning": null,
+                        "store": false,
+                        "temperature": 1.0,
+                        "text": {"format": {"type": "text"}},
+                        "tool_choice": "auto",
+                        "tools": [],
+                        "top_p": 1.0,
+                        "truncation": "disabled",
+                        "usage": null,
+                        "metadata": null
+                    }
+                }),
+            ),
+            (
+                "response.output_item.added",
+                json!({
+                    "type": "response.output_item.added",
+                    "sequence_number": 1,
+                    "output_index": 0,
+                    "item": {
+                        "id": message_id,
+                        "type": "message",
+                        "role": "assistant",
+                        "status": "in_progress",
+                        "content": []
+                    }
+                }),
+            ),
+            (
+                "response.output_text.delta",
+                json!({
+                    "type": "response.output_text.delta",
+                    "sequence_number": 2,
+                    "item_id": message_id,
+                    "output_index": 0,
+                    "content_index": 0,
+                    "delta": "Once upon a streamed unicorn adventure.",
+                    "logprobs": []
+                }),
+            ),
+            (
+                "response.output_text.done",
+                json!({
+                    "type": "response.output_text.done",
+                    "sequence_number": 3,
+                    "item_id": message_id,
+                    "output_index": 0,
+                    "content_index": 0,
+                    "text": final_text,
+                    "logprobs": []
+                }),
+            ),
+            (
+                "response.output_item.done",
+                json!({
+                    "type": "response.output_item.done",
+                    "sequence_number": 4,
+                    "output_index": 0,
+                    "item": {
+                        "id": message_id,
+                        "type": "message",
+                        "role": "assistant",
+                        "status": "completed",
+                        "content": [{
+                            "type": "output_text",
+                            "text": final_text,
+                            "annotations": [],
+                            "logprobs": []
+                        }]
+                    }
+                }),
+            ),
+            (
+                "response.completed",
+                json!({
+                    "type": "response.completed",
+                    "sequence_number": 5,
+                    "response": {
+                        "id": response_id,
+                        "object": "response",
+                        "created_at": 1_700_000_500,
+                        "status": "completed",
+                        "model": "",
+                        "output": [{
+                            "id": message_id,
+                            "type": "message",
+                            "role": "assistant",
+                            "status": "completed",
+                            "content": [{
+                                "type": "output_text",
+                                "text": final_text,
+                                "annotations": [],
+                                "logprobs": []
+                            }]
+                        }],
+                        "parallel_tool_calls": true,
+                        "previous_response_id": null,
+                        "reasoning": null,
+                        "store": false,
+                        "temperature": 1.0,
+                        "text": {"format": {"type": "text"}},
+                        "tool_choice": "auto",
+                        "tools": [],
+                        "top_p": 1.0,
+                        "truncation": "disabled",
+                        "usage": {
+                            "input_tokens": 10,
+                            "input_tokens_details": {"cached_tokens": 0},
+                            "output_tokens": 20,
+                            "output_tokens_details": {"reasoning_tokens": 5},
+                            "total_tokens": 30
+                        },
+                        "metadata": null,
+                        "instructions": null,
+                        "user": null
+                    }
+                }),
+            ),
+        ];
+
+        let sse_payload = events
+            .into_iter()
+            .map(|(event, data)| format!("event: {}\ndata: {}\n\n", event, data))
+            .collect::<String>();
+
+        Response::builder()
+            .status(StatusCode::OK)
+            .header("content-type", "text/event-stream")
+            .body(Body::from(sse_payload))
+            .unwrap()
+    });
+
+    let app = Router::new().route("/v1/responses", sse_handler);
+
+    let server = tokio::spawn(async move {
+        axum::serve(listener, app).await.unwrap();
+    });
+
+    let base_url = format!("http://{}", addr);
+
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    crate::common::test_app::register_external_worker(&ctx, &base_url, Some(vec!["gpt-5-nano"]));
+    let router = OpenAIRouter::new(&ctx).await.unwrap();
+
+    // Get storage from context and seed a previous response
+    let storage = ctx.response_storage.clone();
+    let mut previous = StoredResponse::new(None);
+    previous.id = ResponseId::from("resp_prev_chain");
+    previous.input = serde_json::json!("Earlier bedtime question");
+    previous.output = serde_json::json!("Earlier answer");
+    storage.store_response(previous).await.unwrap();
+
+    let mut metadata = HashMap::new();
+    metadata.insert("topic".to_string(), json!("unicorns"));
+
+    let request = ResponsesRequest {
+        model: "gpt-5-nano".to_string(),
+        input: ResponseInput::Text("Tell me a bedtime story.".to_string()),
+        instructions: Some("Be kind".to_string()),
+        metadata: Some(metadata),
+        previous_response_id: Some("resp_prev_chain".to_string()),
+        store: Some(true),
+        stream: Some(true),
+        ..Default::default()
+    };
+
+    let response = router.route_responses(None, &request, None).await;
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let headers = response.headers();
+    let ct = headers
+        .get("content-type")
+        .unwrap()
+        .to_str()
+        .unwrap()
+        .to_ascii_lowercase();
+    assert!(ct.contains("text/event-stream"));
+
+    let response_body = axum::body::to_bytes(response.into_body(), usize::MAX)
+        .await
+        .unwrap();
+    let body_text = String::from_utf8(response_body.to_vec()).unwrap();
+    assert!(body_text.contains("response.completed"));
+    assert!(body_text.contains("Once upon a streamed unicorn adventure."));
+
+    // Wait for the storage task to persist the streaming response.
+    let target_id = ResponseId::from("resp_stream_123");
+    let stored = loop {
+        if let Some(resp) = storage.get_response(&target_id).await.unwrap() {
+            break resp;
+        }
+        sleep(Duration::from_millis(10)).await;
+    };
+
+    // Input is now stored as a JSON array of items
+    assert!(stored.input.is_array());
+    let input_items = stored.input.as_array().unwrap();
+    assert_eq!(input_items.len(), 1);
+    assert_eq!(input_items[0]["type"], "message");
+    assert_eq!(input_items[0]["role"], "user");
+    assert_eq!(
+        input_items[0]["content"][0]["text"],
+        "Tell me a bedtime story."
+    );
+
+    // Output is now stored as a JSON array of items
+    assert!(stored.output.is_array());
+    let output_items = stored.output.as_array().unwrap();
+    assert_eq!(output_items.len(), 1);
+    assert_eq!(
+        output_items[0]["content"][0]["text"],
+        "Once upon a streamed unicorn adventure."
+    );
+    assert_eq!(
+        stored
+            .previous_response_id
+            .as_ref()
+            .expect("previous_response_id missing")
+            .0,
+        "resp_prev_chain"
+    );
+    assert_eq!(stored.metadata.get("topic"), Some(&json!("unicorns")));
+    assert_eq!(stored.instructions.as_deref(), Some("Be kind"));
+    assert_eq!(stored.model.as_deref(), Some("gpt-5-nano"));
+    assert_eq!(stored.safety_identifier, None);
+    assert_eq!(stored.raw_response["store"], json!(true));
+    assert_eq!(
+        stored.raw_response["previous_response_id"].as_str(),
+        Some("resp_prev_chain")
+    );
+    assert_eq!(stored.raw_response["metadata"]["topic"], json!("unicorns"));
+    assert_eq!(
+        stored.raw_response["instructions"].as_str(),
+        Some("Be kind")
+    );
+
+    server.abort();
+}
+
+/// Test router factory with OpenAI routing mode
+#[tokio::test]
+async fn test_router_factory_openai_mode() {
+    let routing_mode = RoutingMode::OpenAI {
+        worker_urls: vec!["https://api.openai.com".to_string()],
+    };
+
+    let router_config = RouterConfig::new(routing_mode, smg::config::PolicyConfig::Random);
+
+    let app_context = crate::common::create_test_context(router_config).await;
+
+    let router = smg::routers::RouterFactory::create_router(&app_context).await;
+    assert!(
+        router.is_ok(),
+        "Router factory should create OpenAI router successfully"
+    );
+
+    let router = router.unwrap();
+    assert_eq!(router.router_type(), "openai");
+}
+
+/// Test that unsupported endpoints return proper error codes
+#[tokio::test]
+async fn test_unsupported_endpoints() {
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    crate::common::test_app::register_external_worker(&ctx, "https://api.openai.com", None);
+    let router = OpenAIRouter::new(&ctx).await.unwrap();
+
+    let generate_request = GenerateRequest {
+        text: Some("Hello world".to_string()),
+        model: None,
+        input_ids: None,
+        input_embeds: None,
+        image_data: None,
+        video_data: None,
+        audio_data: None,
+        sampling_params: None,
+        return_logprob: Some(false),
+        logprob_start_len: None,
+        top_logprobs_num: None,
+        token_ids_logprob: None,
+        return_text_in_logprobs: false,
+        stream: false,
+        log_metrics: true,
+        return_hidden_states: false,
+        modalities: None,
+        session_params: None,
+        lora_path: None,
+        lora_id: None,
+        custom_logit_processor: None,
+        bootstrap_host: None,
+        bootstrap_port: None,
+        bootstrap_room: None,
+        bootstrap_pair_key: None,
+        data_parallel_rank: None,
+        background: false,
+        conversation_id: None,
+        priority: None,
+        extra_key: None,
+        no_logs: false,
+        custom_labels: None,
+        return_bytes: false,
+        return_entropy: false,
+        rid: None,
+    };
+
+    let response = router.route_generate(None, &generate_request, None).await;
+    assert_eq!(response.status(), StatusCode::NOT_IMPLEMENTED);
+
+    let completion_request = create_minimal_completion_request();
+    let response = router
+        .route_completion(None, &completion_request, None)
+        .await;
+    assert_eq!(response.status(), StatusCode::NOT_IMPLEMENTED);
+}
+
+/// Test chat completion with mock OpenAI server
+#[tokio::test]
+async fn test_openai_router_chat_completion_with_mock() {
+    // Start a mock OpenAI server
+    let mock_server = MockOpenAIServer::new().await;
+    let base_url = mock_server.base_url();
+
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    // Register the mock server worker and create router
+    crate::common::test_app::register_external_worker(&ctx, &base_url, None);
+    let router = OpenAIRouter::new(&ctx).await.unwrap();
+
+    // Create a minimal chat completion request
+    let mut chat_request = create_minimal_chat_request();
+    chat_request.messages = vec![ChatMessage::User {
+        content: MessageContent::Text("Hello, how are you?".to_string()),
+        name: None,
+    }];
+    chat_request.temperature = Some(0.7);
+
+    // Route the request
+    let response = router.route_chat(None, &chat_request, None).await;
+
+    // Should get a successful response from mock server
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let (_, body) = response.into_parts();
+    let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap();
+    let body_str = String::from_utf8(body_bytes.to_vec()).unwrap();
+    let chat_response: serde_json::Value = serde_json::from_str(&body_str).unwrap();
+
+    assert_eq!(chat_response["object"], "chat.completion");
+    assert_eq!(chat_response["model"], "gpt-3.5-turbo");
+    assert!(!chat_response["choices"].as_array().unwrap().is_empty());
+}
+
+/// Test full E2E flow with Axum server
+#[tokio::test]
+async fn test_openai_e2e_with_server() {
+    // Start mock OpenAI server
+    let mock_server = MockOpenAIServer::new().await;
+    let base_url = mock_server.base_url();
+
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    // Register the mock server worker and create router
+    crate::common::test_app::register_external_worker(&ctx, &base_url, None);
+    let router = OpenAIRouter::new(&ctx).await.unwrap();
+
+    // Create Axum app with chat completions endpoint
+    let app = Router::new().route(
+        "/v1/chat/completions",
+        post({
+            let router = Arc::new(router);
+            move |req: Request<Body>| {
+                let router = router.clone();
+                async move {
+                    let (parts, body) = req.into_parts();
+                    let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap();
+                    let body_str = String::from_utf8(body_bytes.to_vec()).unwrap();
+
+                    let chat_request: ChatCompletionRequest =
+                        serde_json::from_str(&body_str).unwrap();
+
+                    router
+                        .route_chat(Some(&parts.headers), &chat_request, None)
+                        .await
+                }
+            }
+        }),
+    );
+
+    // Make a request to the server
+    let request = Request::builder()
+        .method(Method::POST)
+        .uri("/v1/chat/completions")
+        .header("content-type", "application/json")
+        .body(Body::from(
+            json!({
+                "model": "gpt-3.5-turbo",
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "Hello, world!"
+                    }
+                ],
+                "max_tokens": 100
+            })
+            .to_string(),
+        ))
+        .unwrap();
+
+    let response = app.oneshot(request).await.unwrap();
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let body = axum::body::to_bytes(response.into_body(), usize::MAX)
+        .await
+        .unwrap();
+    let response_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
+
+    assert_eq!(response_json["object"], "chat.completion");
+    assert_eq!(response_json["model"], "gpt-3.5-turbo");
+    assert!(!response_json["choices"].as_array().unwrap().is_empty());
+}
+
+/// Test streaming chat completions pass-through with mock server
+#[tokio::test]
+async fn test_openai_router_chat_streaming_with_mock() {
+    let mock_server = MockOpenAIServer::new().await;
+    let base_url = mock_server.base_url();
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    crate::common::test_app::register_external_worker(&ctx, &base_url, None);
+    let router = OpenAIRouter::new(&ctx).await.unwrap();
+
+    // Build a streaming chat request
+    let val = json!({
+        "model": "gpt-3.5-turbo",
+        "messages": [
+            {"role": "user", "content": "Hello"}
+        ],
+        "max_tokens": 10,
+        "stream": true
+    });
+    let chat_request: ChatCompletionRequest = serde_json::from_value(val).unwrap();
+
+    let response = router.route_chat(None, &chat_request, None).await;
+    assert_eq!(response.status(), StatusCode::OK);
+
+    // Should be SSE
+    let headers = response.headers();
+    let ct = headers
+        .get("content-type")
+        .unwrap()
+        .to_str()
+        .unwrap()
+        .to_ascii_lowercase();
+    assert!(ct.contains("text/event-stream"));
+
+    // Read entire stream body and assert chunks + DONE
+    let body = axum::body::to_bytes(response.into_body(), usize::MAX)
+        .await
+        .unwrap();
+    let text = String::from_utf8(body.to_vec()).unwrap();
+    assert!(text.contains("chat.completion.chunk"));
+    assert!(text.contains("[DONE]"));
+}
+
+/// Test circuit breaker functionality
+#[tokio::test]
+async fn test_openai_router_circuit_breaker() {
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    crate::common::test_app::register_external_worker(
+        &ctx,
+        "http://invalid-url-that-will-fail",
+        None,
+    );
+    let router = OpenAIRouter::new(&ctx).await.unwrap();
+
+    let chat_request = create_minimal_chat_request();
+
+    // First few requests should fail and record failures
+    for _ in 0..3 {
+        let response = router.route_chat(None, &chat_request, None).await;
+        // Should get either an error or circuit breaker response
+        assert!(
+            response.status() == StatusCode::INTERNAL_SERVER_ERROR
+                || response.status() == StatusCode::SERVICE_UNAVAILABLE
+        );
+    }
+}
+
+/// Test that /v1/models returns models from registered workers' ModelCards
+///
+/// With the new worker-based design, models are returned from the WorkerRegistry
+/// and don't require calling external APIs. Auth headers are used for routing
+/// requests to workers, not for the models endpoint.
+#[tokio::test]
+async fn test_openai_router_models_from_registry() {
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    // Register a worker with the default model
+    crate::common::test_app::register_external_worker(&ctx, "https://api.example.com", None);
+    let router = OpenAIRouter::new(&ctx).await.unwrap();
+
+    // Get models - should return the registered model
+    let req = Request::builder()
+        .method(Method::GET)
+        .uri("/models")
+        .body(Body::empty())
+        .unwrap();
+
+    let response = router.get_models(req).await;
+    assert_eq!(response.status(), StatusCode::OK);
+    let (_, body) = response.into_parts();
+    let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap();
+    let body_str = String::from_utf8(body_bytes.to_vec()).unwrap();
+    let models: serde_json::Value = serde_json::from_str(&body_str).unwrap();
+    assert_eq!(models["object"], "list");
+
+    // Should have the default model (gpt-3.5-turbo)
+    let data = models["data"].as_array().unwrap();
+    assert_eq!(data.len(), 1);
+    assert_eq!(data[0]["id"], "gpt-3.5-turbo");
+}
+
+#[test]
+fn oracle_config_validation_requires_config_when_enabled() {
+    let config = RouterConfig::builder()
+        .openai_mode(vec!["https://api.openai.com".to_string()])
+        .history_backend(HistoryBackend::Oracle)
+        .build_unchecked();
+
+    let err =
+        ConfigValidator::validate(&config).expect_err("config should fail without oracle details");
+
+    match err {
+        ConfigError::MissingRequired { field } => {
+            assert_eq!(field, "oracle");
+        }
+        other => panic!("unexpected error: {:?}", other),
+    }
+}
+
+#[test]
+fn oracle_config_validation_accepts_dsn_only() {
+    let config = RouterConfig::builder()
+        .openai_mode(vec!["https://api.openai.com".to_string()])
+        .oracle_history(OracleConfig {
+            wallet_path: None,
+            connect_descriptor: "tcps://db.example.com:1522/service".to_string(),
+            username: "scott".to_string(),
+            password: "tiger".to_string(),
+            pool_min: 1,
+            pool_max: 4,
+            pool_timeout_secs: 30,
+        })
+        .build_unchecked();
+
+    ConfigValidator::validate(&config).expect("dsn-based config should validate");
+}
+
+#[test]
+fn oracle_config_validation_accepts_wallet_alias() {
+    let config = RouterConfig::builder()
+        .openai_mode(vec!["https://api.openai.com".to_string()])
+        .oracle_history(OracleConfig {
+            wallet_path: Some("/etc/sglang/oracle-wallet".to_string()),
+            connect_descriptor: "db_low".to_string(),
+            username: "app_user".to_string(),
+            password: "secret".to_string(),
+            pool_min: 1,
+            pool_max: 8,
+            pool_timeout_secs: 45,
+        })
+        .build_unchecked();
+
+    ConfigValidator::validate(&config).expect("wallet-based config should validate");
+}

sgl-model-gateway/tests/routing/test_pd_routing.rs

@@ -0,0 +1,927 @@
+#[cfg(test)]
+mod pd_routing_unit_tests {
+    use serde_json::json;
+    use smg::{
+        app_context::AppContext,
+        config::{PolicyConfig, RouterConfig, RoutingMode},
+        core::{BasicWorkerBuilder, Worker, WorkerType},
+        routers::{http::pd_types::PDSelectionPolicy, RouterFactory},
+        tokenizer::registry::TokenizerRegistry,
+    };
+
+    #[derive(Debug)]
+    struct PDRequest {
+        pub is_stream: bool,
+        pub batch_size: Option<usize>,
+    }
+
+    impl PDRequest {
+        pub fn from_json(json: &serde_json::Value) -> Self {
+            let is_stream = json
+                .get("stream")
+                .and_then(|v| v.as_bool())
+                .unwrap_or(false);
+
+            let batch_size = if let Some(text) = json.get("text") {
+                text.as_array().map(|arr| arr.len())
+            } else if let Some(input_ids) = json.get("input_ids") {
+                input_ids.as_array().map(|arr| arr.len())
+            } else {
+                None
+            };
+
+            PDRequest {
+                is_stream,
+                batch_size,
+            }
+        }
+    }
+
+    #[test]
+    fn test_worker_types() {
+        use smg::core::{BasicWorkerBuilder, Worker, WorkerType};
+
+        let prefill_worker: Box<dyn Worker> = Box::new(
+            BasicWorkerBuilder::new("http://prefill:8080")
+                .worker_type(WorkerType::Prefill {
+                    bootstrap_port: Some(9000),
+                })
+                .api_key("test_api_key")
+                .build(),
+        );
+        assert_eq!(prefill_worker.url(), "http://prefill:8080");
+        match prefill_worker.worker_type() {
+            WorkerType::Prefill { bootstrap_port } => {
+                assert_eq!(*bootstrap_port, Some(9000));
+            }
+            _ => panic!("Expected Prefill worker type"),
+        }
+
+        let decode_worker: Box<dyn Worker> = Box::new(
+            BasicWorkerBuilder::new("http://decode:8080")
+                .worker_type(WorkerType::Decode)
+                .api_key("test_api_key")
+                .build(),
+        );
+        assert_eq!(decode_worker.url(), "http://decode:8080");
+        match decode_worker.worker_type() {
+            WorkerType::Decode => (),
+            _ => panic!("Expected Decode worker type"),
+        }
+
+        let regular_worker: Box<dyn Worker> = Box::new(
+            BasicWorkerBuilder::new("http://regular:8080")
+                .worker_type(WorkerType::Regular)
+                .api_key("test_api_key")
+                .build(),
+        );
+        assert_eq!(regular_worker.url(), "http://regular:8080");
+        match regular_worker.worker_type() {
+            WorkerType::Regular => (),
+            _ => panic!("Expected Regular worker type"),
+        }
+    }
+
+    #[test]
+    fn test_pd_selection_policies() {
+        // Note: These policies are only used when pd_disaggregation=true
+        let policies = vec![
+            PDSelectionPolicy::Random,
+            PDSelectionPolicy::PowerOfTwo,
+            PDSelectionPolicy::CacheAware {
+                cache_threshold: 0.5,
+                balance_abs_threshold: 32,
+                balance_rel_threshold: 1.1,
+            },
+            PDSelectionPolicy::Bucket {
+                balance_abs_threshold: 32,
+                balance_rel_threshold: 1.1,
+                bucket_adjust_interval_secs: 5,
+            },
+        ];
+
+        for policy in policies {
+            match &policy {
+                PDSelectionPolicy::Random => {
+                    assert!(matches!(policy, PDSelectionPolicy::Random));
+                }
+                PDSelectionPolicy::PowerOfTwo => {
+                    assert!(matches!(policy, PDSelectionPolicy::PowerOfTwo));
+                }
+                PDSelectionPolicy::CacheAware {
+                    cache_threshold, ..
+                } => {
+                    assert!(*cache_threshold >= 0.0 && *cache_threshold <= 1.0);
+                }
+                PDSelectionPolicy::Bucket {
+                    balance_rel_threshold,
+                    ..
+                } => {
+                    assert!(*balance_rel_threshold >= 1.0);
+                }
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_pd_router_configuration() {
+        // In the new structure, RoutingMode and PolicyConfig are separate
+        let test_cases = vec![
+            (
+                RoutingMode::PrefillDecode {
+                    prefill_urls: vec![
+                        ("http://prefill1:8080".to_string(), Some(9000)),
+                        ("http://prefill2:8080".to_string(), None),
+                    ],
+                    decode_urls: vec![
+                        "http://decode1:8080".to_string(),
+                        "http://decode2:8080".to_string(),
+                    ],
+                    prefill_policy: None,
+                    decode_policy: None,
+                },
+                PolicyConfig::Random,
+            ),
+            (
+                RoutingMode::PrefillDecode {
+                    prefill_urls: vec![("http://prefill:8080".to_string(), Some(9000))],
+                    decode_urls: vec!["http://decode:8080".to_string()],
+                    prefill_policy: None,
+                    decode_policy: None,
+                },
+                PolicyConfig::PowerOfTwo {
+                    load_check_interval_secs: 5,
+                },
+            ),
+            (
+                RoutingMode::PrefillDecode {
+                    prefill_urls: vec![
+                        ("http://p1:8080".to_string(), Some(9000)),
+                        ("http://p2:8080".to_string(), Some(9001)),
+                        ("http://p3:8080".to_string(), Some(9002)),
+                    ],
+                    decode_urls: vec!["http://d1:8080".to_string(), "http://d2:8080".to_string()],
+                    prefill_policy: None,
+                    decode_policy: None,
+                },
+                PolicyConfig::CacheAware {
+                    cache_threshold: 0.7,
+                    balance_abs_threshold: 20,
+                    balance_rel_threshold: 1.2,
+                    eviction_interval_secs: 60,
+                    max_tree_size: 1000000,
+                },
+            ),
+            (
+                RoutingMode::PrefillDecode {
+                    prefill_urls: vec![
+                        ("http://p1:8080".to_string(), Some(9000)),
+                        ("http://p2:8080".to_string(), Some(9001)),
+                        ("http://p3:8080".to_string(), Some(9002)),
+                    ],
+                    decode_urls: vec!["http://d1:8080".to_string(), "http://d2:8080".to_string()],
+                    prefill_policy: None,
+                    decode_policy: None,
+                },
+                PolicyConfig::Bucket {
+                    balance_abs_threshold: 20,
+                    balance_rel_threshold: 1.2,
+                    bucket_adjust_interval_secs: 5,
+                },
+            ),
+        ];
+
+        for (mode, policy) in test_cases {
+            let config = match mode {
+                RoutingMode::PrefillDecode {
+                    prefill_urls,
+                    decode_urls,
+                    ..
+                } => RouterConfig::builder()
+                    .prefill_decode_mode(prefill_urls, decode_urls)
+                    .policy(policy)
+                    .host("127.0.0.1")
+                    .port(3001)
+                    .max_payload_size(1024 * 1024)
+                    .request_timeout_secs(60)
+                    .worker_startup_timeout_secs(10)
+                    .worker_startup_check_interval_secs(1)
+                    .max_concurrent_requests(64)
+                    .queue_timeout_secs(60)
+                    .build_unchecked(),
+                _ => panic!("Expected PrefillDecode mode"),
+            };
+
+            let app_context = {
+                use std::sync::{Arc, OnceLock};
+
+                use smg::{
+                    core::{LoadMonitor, WorkerRegistry},
+                    data_connector::{
+                        MemoryConversationItemStorage, MemoryConversationStorage,
+                        MemoryResponseStorage,
+                    },
+                    middleware::TokenBucket,
+                    policies::PolicyRegistry,
+                };
+
+                let client = reqwest::Client::new();
+
+                // Initialize rate limiter
+                let rate_limiter = Some(Arc::new(TokenBucket::new(64, 64)));
+
+                // Initialize registries
+                let worker_registry = Arc::new(WorkerRegistry::new());
+                let policy_registry = Arc::new(PolicyRegistry::new(config.policy.clone()));
+
+                // Initialize storage backends
+                let response_storage = Arc::new(MemoryResponseStorage::new());
+                let conversation_storage = Arc::new(MemoryConversationStorage::new());
+                let conversation_item_storage = Arc::new(MemoryConversationItemStorage::new());
+
+                // Initialize load monitor
+                let load_monitor = Some(Arc::new(LoadMonitor::new(
+                    worker_registry.clone(),
+                    policy_registry.clone(),
+                    client.clone(),
+                    config.worker_startup_check_interval_secs,
+                )));
+
+                // Create empty OnceLock for worker job queue, workflow engine, and mcp manager
+                let worker_job_queue = Arc::new(OnceLock::new());
+                let workflow_engine = Arc::new(OnceLock::new());
+                let mcp_manager = Arc::new(OnceLock::new());
+
+                Arc::new(
+                    AppContext::builder()
+                        .router_config(config)
+                        .client(client)
+                        .rate_limiter(rate_limiter)
+                        .tokenizer_registry(Arc::new(TokenizerRegistry::new())) // tokenizer
+                        .reasoning_parser_factory(None) // reasoning_parser_factory
+                        .tool_parser_factory(None) // tool_parser_factory
+                        .worker_registry(worker_registry)
+                        .policy_registry(policy_registry)
+                        .response_storage(response_storage)
+                        .conversation_storage(conversation_storage)
+                        .conversation_item_storage(conversation_item_storage)
+                        .load_monitor(load_monitor)
+                        .worker_job_queue(worker_job_queue)
+                        .workflow_engine(workflow_engine)
+                        .mcp_manager(mcp_manager)
+                        .build()
+                        .unwrap(),
+                )
+            };
+            let result = RouterFactory::create_router(&app_context).await;
+            assert!(
+                result.is_ok(),
+                "Router creation should succeed with empty worker"
+            );
+
+            let stats = app_context.worker_registry.stats();
+            assert_eq!(
+                stats.total_workers, 0,
+                "No workers should be registered without initialization"
+            );
+        }
+    }
+
+    #[test]
+    fn test_pd_request_from_json() {
+        let single_json = json!({
+            "text": "Hello world",
+            "stream": false,
+            "temperature": 0.7,
+            "max_tokens": 100
+        });
+
+        let pd_req = PDRequest::from_json(&single_json);
+        assert!(!pd_req.is_stream);
+        assert_eq!(pd_req.batch_size, None);
+
+        let batch_json = json!({
+            "text": ["Hello", "World", "Test"],
+            "stream": true,
+            "temperature": 0.5
+        });
+
+        let pd_req = PDRequest::from_json(&batch_json);
+        assert!(pd_req.is_stream);
+        assert_eq!(pd_req.batch_size, Some(3));
+
+        let ids_json = json!({
+            "input_ids": [[1, 2, 3], [4, 5, 6]],
+            "stream": false
+        });
+
+        let pd_req = PDRequest::from_json(&ids_json);
+        assert!(!pd_req.is_stream);
+        assert_eq!(pd_req.batch_size, Some(2));
+
+        let chat_json = json!({
+            "messages": [
+                {"role": "system", "content": "You are a helpful assistant"},
+                {"role": "user", "content": "Hello"}
+            ],
+            "stream": true
+        });
+
+        let pd_req = PDRequest::from_json(&chat_json);
+        assert!(pd_req.is_stream);
+        assert_eq!(pd_req.batch_size, None);
+    }
+
+    #[test]
+    fn test_bootstrap_injection_simulation() {
+        // Since we can't test the actual inject_bootstrap_fields function here
+        // (it's private in the router module), we'll test the expected behavior
+
+        let mut single_json = json!({
+            "text": "Hello world",
+            "stream": false,
+            "temperature": 0.7
+        });
+
+        let prefill_worker: Box<dyn Worker> = Box::new(
+            BasicWorkerBuilder::new("http://prefill1:8080")
+                .worker_type(WorkerType::Prefill {
+                    bootstrap_port: Some(9000),
+                })
+                .api_key("test_api_key")
+                .build(),
+        );
+
+        let bootstrap_port = match prefill_worker.worker_type() {
+            WorkerType::Prefill { bootstrap_port } => bootstrap_port,
+            _ => &None,
+        };
+
+        single_json["bootstrap_host"] = json!(prefill_worker.bootstrap_host());
+        single_json["bootstrap_port"] = json!(bootstrap_port);
+        single_json["bootstrap_room"] = json!(12345u64); // Random room ID
+
+        assert_eq!(single_json["bootstrap_host"], "prefill1");
+        assert_eq!(single_json["bootstrap_port"], json!(Some(9000)));
+        assert!(single_json["bootstrap_room"].is_u64());
+        assert_eq!(single_json["temperature"], 0.7); // Original field preserved
+
+        let mut batch_json = json!({
+            "text": ["Hello", "World", "Test"],
+            "stream": true
+        });
+
+        let batch_size = 3;
+        let hostname = prefill_worker.bootstrap_host();
+        batch_json["bootstrap_host"] = json!(vec![hostname; batch_size]);
+        batch_json["bootstrap_port"] = json!(vec![bootstrap_port; batch_size]);
+        batch_json["bootstrap_room"] = json!(vec![111u64, 222u64, 333u64]);
+
+        assert!(batch_json["bootstrap_host"].is_array());
+        assert_eq!(
+            batch_json["bootstrap_host"].as_array().unwrap().len(),
+            batch_size
+        );
+        assert!(batch_json["bootstrap_port"].is_array());
+        assert!(batch_json["bootstrap_room"].is_array());
+        assert_eq!(batch_json["stream"], true); // Original field preserved
+    }
+
+    #[test]
+    fn test_request_serialization() {
+        let request = json!({
+            "text": "Test prompt",
+            "stream": false,
+            "temperature": 0.7,
+            "max_tokens": 100,
+            "top_p": 0.9,
+            "frequency_penalty": 0.5,
+            "bootstrap_host": "prefill1",
+            "bootstrap_port": 9000,
+            "bootstrap_room": 12345u64
+        });
+
+        let bytes = serde_json::to_vec(&request).unwrap();
+
+        let parsed: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
+
+        assert_eq!(parsed["text"], "Test prompt");
+        assert_eq!(parsed["stream"], false);
+        assert_eq!(parsed["temperature"], 0.7);
+        assert_eq!(parsed["max_tokens"], 100);
+        assert_eq!(parsed["bootstrap_host"], "prefill1");
+        assert_eq!(parsed["bootstrap_port"], 9000);
+        assert_eq!(parsed["bootstrap_room"], 12345);
+    }
+
+    #[test]
+    fn test_pd_request_edge_cases() {
+        let empty_json = json!({});
+        let pd_req = PDRequest::from_json(&empty_json);
+        assert!(!pd_req.is_stream);
+        assert_eq!(pd_req.batch_size, None);
+
+        let stream_only = json!({
+            "stream": true
+        });
+        let pd_req = PDRequest::from_json(&stream_only);
+        assert!(pd_req.is_stream);
+        assert_eq!(pd_req.batch_size, None);
+
+        let empty_batch = json!({
+            "text": []
+        });
+        let pd_req = PDRequest::from_json(&empty_batch);
+        assert_eq!(pd_req.batch_size, Some(0));
+
+        let non_array_text = json!({
+            "text": "single string"
+        });
+        let pd_req = PDRequest::from_json(&non_array_text);
+        assert_eq!(pd_req.batch_size, None);
+    }
+
+    #[tokio::test]
+    async fn test_background_load_monitoring() {
+        use std::collections::HashMap;
+
+        use tokio::sync::watch;
+
+        let (tx, rx) = watch::channel(HashMap::new());
+
+        let mut loads = HashMap::new();
+        loads.insert("http://prefill1:8080".to_string(), 10);
+        loads.insert("http://prefill2:8080".to_string(), 20);
+        loads.insert("http://decode1:8080".to_string(), 5);
+        loads.insert("http://decode2:8080".to_string(), 15);
+
+        tx.send(loads.clone()).unwrap();
+
+        let received_loads = rx.borrow();
+        assert_eq!(received_loads.get("http://prefill1:8080"), Some(&10));
+        assert_eq!(received_loads.get("http://prefill2:8080"), Some(&20));
+        assert_eq!(received_loads.get("http://decode1:8080"), Some(&5));
+        assert_eq!(received_loads.get("http://decode2:8080"), Some(&15));
+    }
+
+    #[test]
+    fn test_load_monitoring_configuration() {
+        let policies = vec![
+            (PDSelectionPolicy::Random, false),
+            (PDSelectionPolicy::PowerOfTwo, true),
+            (
+                PDSelectionPolicy::CacheAware {
+                    cache_threshold: 0.5,
+                    balance_abs_threshold: 32,
+                    balance_rel_threshold: 1.1,
+                },
+                false,
+            ),
+        ];
+
+        for (policy, should_monitor) in policies {
+            match policy {
+                PDSelectionPolicy::PowerOfTwo => assert!(should_monitor),
+                _ => assert!(!should_monitor),
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_watch_channel_behavior() {
+        use std::collections::HashMap;
+
+        use tokio::sync::watch;
+
+        let (tx, rx1) = watch::channel(HashMap::new());
+        let rx2 = rx1.clone();
+
+        assert!(rx1.borrow().is_empty());
+        assert!(rx2.borrow().is_empty());
+
+        let mut loads = HashMap::new();
+        loads.insert("worker1".to_string(), 10);
+        tx.send(loads.clone()).unwrap();
+
+        assert_eq!(rx1.borrow().get("worker1"), Some(&10));
+        assert_eq!(rx2.borrow().get("worker1"), Some(&10));
+
+        loads.insert("worker1".to_string(), 20);
+        loads.insert("worker2".to_string(), 30);
+        tx.send(loads).unwrap();
+
+        assert_eq!(rx1.borrow().get("worker1"), Some(&20));
+        assert_eq!(rx2.borrow().get("worker2"), Some(&30));
+    }
+
+    #[test]
+    fn test_generate_request_formats() {
+        // Based on bench_one_batch_server.py request patterns
+
+        let batch_request = json!({
+            "input_ids": [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]],
+            "sampling_params": {
+                "temperature": 0.0,
+                "max_new_tokens": 16,
+                "ignore_eos": true,
+            },
+            "return_logprob": false,
+            "stream": true
+        });
+
+        let pd_req = PDRequest::from_json(&batch_request);
+        assert!(pd_req.is_stream);
+        assert_eq!(pd_req.batch_size, Some(3));
+
+        let logprob_request = json!({
+            "input_ids": [[1, 2, 3]],
+            "sampling_params": {
+                "temperature": 0.7,
+                "max_new_tokens": 8,
+            },
+            "return_logprob": true,
+            "stream": false
+        });
+
+        assert_eq!(logprob_request["return_logprob"], true);
+        assert_eq!(logprob_request["stream"], false);
+
+        let batch_sizes = vec![1, 16, 64]; // From bench_one_batch_server.py
+        for bs in batch_sizes {
+            let request = json!({
+                "input_ids": vec![vec![1, 2, 3]; bs],
+                "sampling_params": {
+                    "temperature": 0.0,
+                    "max_new_tokens": 16,
+                },
+                "stream": true
+            });
+
+            let pd_req = PDRequest::from_json(&request);
+            assert_eq!(pd_req.batch_size, Some(bs));
+        }
+    }
+
+    #[test]
+    fn test_sampling_params_handling() {
+        let sampling_params_variations = vec![
+            json!({
+                "temperature": 0.0,
+                "max_new_tokens": 8,
+                "ignore_eos": true
+            }),
+            json!({
+                "temperature": 0.7,
+                "max_new_tokens": 16,
+                "ignore_eos": false,
+                "top_p": 0.9,
+                "frequency_penalty": 0.5
+            }),
+            json!({
+                "temperature": 1.0,
+                "max_new_tokens": 64,
+                "json_schema": "$$ANY$$"  // Structured output
+            }),
+        ];
+
+        for params in sampling_params_variations {
+            let request = json!({
+                "input_ids": [[1, 2, 3]],
+                "sampling_params": params.clone(),
+                "stream": false
+            });
+
+            assert_eq!(request["sampling_params"], params);
+        }
+    }
+
+    #[test]
+    fn test_streaming_response_parsing() {
+        let sse_chunks = ["data: {\"text\":\"Hello\",\"meta_info\":{\"completion_tokens\":1,\"finish_reason\":null}}",
+            "data: {\"text\":\" world\",\"meta_info\":{\"completion_tokens\":2,\"finish_reason\":null}}",
+            "data: {\"text\":\"!\",\"meta_info\":{\"completion_tokens\":3,\"finish_reason\":{\"type\":\"length\"}}}",
+            "data: [DONE]"];
+
+        for chunk in &sse_chunks[..3] {
+            assert!(chunk.starts_with("data: "));
+            let json_str = &chunk[6..]; // Skip "data: "
+            let parsed: serde_json::Value = serde_json::from_str(json_str).unwrap();
+            assert!(parsed["meta_info"]["completion_tokens"].is_u64());
+        }
+
+        assert_eq!(sse_chunks[3], "data: [DONE]");
+    }
+
+    #[test]
+    fn test_ttft_calculation() {
+        let first_token_response = json!({
+            "text": "Hello",
+            "meta_info": {
+                "completion_tokens": 1,
+                "finish_reason": null
+            }
+        });
+
+        // TTFT is calculated when completion_tokens == 1
+        assert_eq!(first_token_response["meta_info"]["completion_tokens"], 1);
+        assert!(first_token_response["meta_info"]["finish_reason"].is_null());
+    }
+
+    #[test]
+    fn test_throughput_metrics() {
+        let batch_size = 16;
+        let input_len = 1024;
+        let output_len = 16;
+        let ttft = 0.5; // seconds
+        let total_latency = 2.0; // seconds
+
+        // Input throughput = batch_size * input_len / ttft
+        let input_throughput = (batch_size as f64) * (input_len as f64) / ttft;
+        assert!((input_throughput - 32768.0).abs() < 0.01);
+
+        // Output throughput = batch_size * output_len / (latency - ttft)
+        let output_throughput = (batch_size as f64) * (output_len as f64) / (total_latency - ttft);
+        assert!((output_throughput - 170.67).abs() < 0.01);
+    }
+
+    #[test]
+    fn test_error_response_handling() {
+        let error_response = json!({
+            "error": "Request has failed. Invalid input format."
+        });
+
+        assert!(error_response.get("error").is_some());
+        assert!(error_response["error"].as_str().unwrap().contains("failed"));
+    }
+
+    #[test]
+    fn test_structured_output_request() {
+        let structured_request = json!({
+            "text": "What is the capital of France? Answer in JSON.",
+            "sampling_params": {
+                "temperature": 0.0,
+                "max_new_tokens": 64,
+                "json_schema": "$$ANY$$"
+            },
+            "stream": false
+        });
+
+        assert_eq!(
+            structured_request["sampling_params"]["json_schema"],
+            "$$ANY$$"
+        );
+    }
+
+    #[test]
+    fn test_bootstrap_injection_with_benchmark_requests() {
+        use smg::core::{BasicWorkerBuilder, Worker, WorkerType};
+
+        let mut benchmark_request = json!({
+            "input_ids": vec![vec![1, 2, 3, 4]; 16], // Batch size 16
+            "sampling_params": {
+                "temperature": 0.0,
+                "max_new_tokens": 8,
+                "ignore_eos": true
+            },
+            "return_logprob": true,
+            "stream": true
+        });
+
+        let prefill_worker: Box<dyn Worker> = Box::new(
+            BasicWorkerBuilder::new("http://prefill:8080")
+                .worker_type(WorkerType::Prefill {
+                    bootstrap_port: Some(9000),
+                })
+                .api_key("test_api_key")
+                .build(),
+        );
+
+        let bootstrap_port = match prefill_worker.worker_type() {
+            WorkerType::Prefill { bootstrap_port } => bootstrap_port,
+            _ => &None,
+        };
+        let batch_size = 16;
+        let hostname = prefill_worker.bootstrap_host();
+
+        benchmark_request["bootstrap_host"] = json!(vec![hostname; batch_size]);
+        benchmark_request["bootstrap_port"] = json!(vec![bootstrap_port; batch_size]);
+        benchmark_request["bootstrap_room"] =
+            json!((0..batch_size).map(|_| 12345u64).collect::<Vec<_>>());
+
+        assert_eq!(
+            benchmark_request["bootstrap_host"]
+                .as_array()
+                .unwrap()
+                .len(),
+            batch_size
+        );
+        assert_eq!(
+            benchmark_request["bootstrap_port"]
+                .as_array()
+                .unwrap()
+                .len(),
+            batch_size
+        );
+        assert_eq!(
+            benchmark_request["bootstrap_room"]
+                .as_array()
+                .unwrap()
+                .len(),
+            batch_size
+        );
+
+        assert_eq!(benchmark_request["return_logprob"], true);
+        assert_eq!(benchmark_request["stream"], true);
+    }
+
+    #[test]
+    fn test_server_info_response_format() {
+        let server_info = json!({
+            "internal_states": [{
+                "avg_spec_accept_length": 3.5,
+                "last_gen_throughput": 2048.5,
+                "load": 16
+            }],
+            "prefill": [
+                {"url": "http://prefill1:8080", "load": 10},
+                {"url": "http://prefill2:8080", "load": 20}
+            ],
+            "decode": [
+                {"url": "http://decode1:8080", "load": 5},
+                {"url": "http://decode2:8080", "load": 15}
+            ]
+        });
+
+        assert!(server_info["internal_states"][0]["avg_spec_accept_length"].is_f64());
+        assert!(server_info["internal_states"][0]["last_gen_throughput"].is_f64());
+        assert!(server_info["prefill"].is_array());
+        assert!(server_info["decode"].is_array());
+    }
+
+    // Comprehensive Endpoint Coverage Test
+
+    #[test]
+    fn test_pd_endpoints_coverage() {
+        // Document all endpoints from Python mini_lb.py and verify implementation status
+        let implemented_endpoints = vec![
+            ("/health", "GET", true),
+            ("/health_generate", "GET", true), // Note: Python uses POST, we use GET
+            ("/get_server_info", "GET", true),
+            ("/v1/models", "GET", true),
+            ("/get_model_info", "GET", true),
+            ("/generate", "POST", true),
+            ("/v1/chat/completions", "POST", true),
+            ("/v1/completions", "POST", true),
+            ("/flush_cache", "POST", true),
+            ("/get_loads", "GET", true),
+            ("/register", "POST", false), // NOT IMPLEMENTED - needs dynamic worker management
+        ];
+
+        let implemented_count = implemented_endpoints
+            .iter()
+            .filter(|(_, _, impl_status)| *impl_status)
+            .count();
+        let total_count = implemented_endpoints.len();
+
+        // We've implemented 10 out of 11 endpoints (register is not needed for Phase 1/2)
+        assert_eq!(implemented_count, 10);
+        assert_eq!(total_count, 11);
+
+        let missing: Vec<_> = implemented_endpoints
+            .iter()
+            .filter(|(_, _, impl_status)| !impl_status)
+            .map(|(endpoint, method, _)| format!("{} {}", method, endpoint))
+            .collect();
+
+        assert_eq!(missing, vec!["POST /register"]);
+    }
+
+    #[test]
+    fn test_large_batch_bootstrap_injection() {
+        // This simulates the bench_one_batch_server.py scenario
+        let large_batch_sizes = vec![1024, 4096, 8192];
+
+        for batch_size in large_batch_sizes {
+            let start = std::time::Instant::now();
+
+            let mut large_batch_request = json!({
+                "input_ids": vec![vec![1, 2, 3, 4]; batch_size],
+                "sampling_params": {
+                    "temperature": 0.0,
+                    "max_new_tokens": 16,
+                },
+                "stream": true
+            });
+
+            let prefill_worker: Box<dyn Worker> = Box::new(
+                BasicWorkerBuilder::new("http://prefill:8080")
+                    .worker_type(WorkerType::Prefill {
+                        bootstrap_port: Some(9000),
+                    })
+                    .api_key("test_api_key")
+                    .build(),
+            );
+
+            let bootstrap_port = match prefill_worker.worker_type() {
+                WorkerType::Prefill { bootstrap_port } => bootstrap_port,
+                _ => &None,
+            };
+            let hostname = prefill_worker.bootstrap_host();
+
+            large_batch_request["bootstrap_host"] = json!(vec![hostname; batch_size]);
+            large_batch_request["bootstrap_port"] = json!(vec![bootstrap_port; batch_size]);
+            large_batch_request["bootstrap_room"] = json!((0..batch_size)
+                .map(|_| rand::random::<u64>())
+                .collect::<Vec<_>>());
+
+            let elapsed = start.elapsed();
+
+            assert_eq!(
+                large_batch_request["bootstrap_host"]
+                    .as_array()
+                    .unwrap()
+                    .len(),
+                batch_size
+            );
+            assert_eq!(
+                large_batch_request["bootstrap_port"]
+                    .as_array()
+                    .unwrap()
+                    .len(),
+                batch_size
+            );
+            assert_eq!(
+                large_batch_request["bootstrap_room"]
+                    .as_array()
+                    .unwrap()
+                    .len(),
+                batch_size
+            );
+
+            // Bootstrap injection should be reasonably fast even for large batches
+            println!(
+                "Bootstrap injection for batch_size {} took {:?}",
+                batch_size, elapsed
+            );
+            assert!(
+                elapsed.as_millis() < 1000,
+                "Bootstrap injection took too long for batch size {}",
+                batch_size
+            );
+        }
+    }
+
+    #[test]
+    fn test_payload_size_calculation() {
+        let test_cases = vec![
+            (1, 1024, 16),   // Small batch
+            (16, 1024, 16),  // Medium batch
+            (64, 1024, 16),  // Large batch
+            (8192, 4096, 5), // Benchmark scenario
+        ];
+
+        for (batch_size, input_len, _output_len) in test_cases {
+            // Estimate payload size (rough calculation)
+            // Each token is ~4 bytes (i32), plus JSON overhead
+            let tokens_size = batch_size * input_len * 4; // 4 bytes per token
+            let json_overhead = batch_size * 100; // ~100 bytes overhead per request
+            let total_size = tokens_size + json_overhead;
+
+            println!(
+                "Batch size: {}, Input len: {}, Estimated payload: {} MB",
+                batch_size,
+                input_len,
+                total_size / (1024 * 1024)
+            );
+
+            // For the benchmark case (8192, 4096), this should be ~134 MB
+            if batch_size == 8192 && input_len == 4096 {
+                assert!(
+                    total_size > 100 * 1024 * 1024,
+                    "Benchmark payload should be > 100MB"
+                );
+                assert!(
+                    total_size < 200 * 1024 * 1024,
+                    "Benchmark payload should be < 200MB"
+                );
+            }
+        }
+    }
+
+    #[test]
+    fn test_policy_type_to_pd_selection_policy_mapping() {
+        let pd_policy_count = 3; // Random, PowerOfTwo, CacheAware
+        assert_eq!(
+            pd_policy_count, 3,
+            "PDSelectionPolicy should have exactly 3 variants"
+        );
+
+        let _random = PDSelectionPolicy::Random;
+        let _po2 = PDSelectionPolicy::PowerOfTwo;
+        let _cache_aware = PDSelectionPolicy::CacheAware {
+            cache_threshold: 0.5,
+            balance_abs_threshold: 32,
+            balance_rel_threshold: 1.1,
+        };
+    }
+}