chore: vendor sglang v0.5.10 snapshot

third_party/sglang/sgl-model-gateway/tests/api/mod.rs

@@ -0,0 +1,7 @@
+//! API endpoint integration tests
+
+mod api_endpoints_test;
+mod parser_endpoints_test;
+mod request_formats_test;
+mod responses_api_test;
+mod streaming_tests;

third_party/sglang/sgl-model-gateway/tests/api/parser_endpoints_test.rs

@@ -0,0 +1,604 @@
+use std::sync::Arc;
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use reqwest::Client;
+use serde_json::json;
+use smg::{
+    app_context::AppContext,
+    config::{RouterConfig, RoutingMode},
+    routers::{RouterFactory, RouterTrait},
+};
+use tower::ServiceExt;
+
+use crate::common::mock_worker::{MockWorker, MockWorkerConfig};
+
+/// Test context that manages mock workers and app
+struct ParserTestContext {
+    workers: Vec<MockWorker>,
+    router: Arc<dyn RouterTrait>,
+    _client: Client,
+    _config: RouterConfig,
+    app_context: Arc<AppContext>,
+}
+
+impl ParserTestContext {
+    async fn new(worker_configs: Vec<MockWorkerConfig>) -> Self {
+        // Create router config with parser support enabled
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .random_policy()
+            .host("127.0.0.1")
+            .port(3003)
+            .max_payload_size(256 * 1024 * 1024)
+            .request_timeout_secs(600)
+            .worker_startup_timeout_secs(1)
+            .worker_startup_check_interval_secs(1)
+            .max_concurrent_requests(64)
+            .queue_timeout_secs(60)
+            .build_unchecked();
+
+        Self::new_with_config(config, worker_configs).await
+    }
+
+    async fn new_with_config(
+        mut config: RouterConfig,
+        worker_configs: Vec<MockWorkerConfig>,
+    ) -> Self {
+        let mut workers = Vec::new();
+        let mut worker_urls = Vec::new();
+
+        // Start mock workers if any
+        for worker_config in worker_configs {
+            let mut worker = MockWorker::new(worker_config);
+            let url = worker.start().await.unwrap();
+            worker_urls.push(url);
+            workers.push(worker);
+        }
+
+        if !workers.is_empty() {
+            tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
+        }
+
+        // Update config with worker URLs if not already set
+        match &mut config.mode {
+            RoutingMode::Regular {
+                worker_urls: ref mut urls,
+            } => {
+                if urls.is_empty() {
+                    *urls = worker_urls.clone();
+                }
+            }
+            RoutingMode::OpenAI {
+                worker_urls: ref mut urls,
+            } => {
+                if urls.is_empty() {
+                    *urls = worker_urls.clone();
+                }
+            }
+            _ => {} // PrefillDecode mode has its own setup
+        }
+
+        let client = Client::builder()
+            .timeout(std::time::Duration::from_secs(config.request_timeout_secs))
+            .build()
+            .unwrap();
+
+        // Create app context with parser factories initialized
+        let app_context = crate::common::create_test_context_with_parsers(config.clone()).await;
+
+        // Create router
+        let router = RouterFactory::create_router(&app_context).await.unwrap();
+        let router = Arc::from(router);
+
+        Self {
+            workers,
+            router,
+            _client: client,
+            _config: config,
+            app_context,
+        }
+    }
+
+    async fn create_app(&self) -> axum::Router {
+        crate::common::test_app::create_test_app_with_context(
+            Arc::clone(&self.router),
+            Arc::clone(&self.app_context),
+        )
+    }
+
+    async fn shutdown(mut self) {
+        for worker in &mut self.workers {
+            worker.stop().await;
+        }
+    }
+}
+
+#[cfg(test)]
+mod parse_function_call_tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_parse_function_call_success() {
+        let ctx = ParserTestContext::new(vec![]).await;
+        let app = ctx.create_app().await;
+
+        let payload = json!({
+            "text": r#"I need to call the weather function <tool_call>{"function_name": "get_weather", "parameters": {"location": "Beijing"}}</tool_call>"#,
+            "tool_call_parser": "json",
+            "tools": [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "description": "Get weather information",
+                        "parameters": {
+                            "type": "object",
+                            "properties": {
+                                "location": {
+                                    "type": "string",
+                                    "description": "The location"
+                                }
+                            },
+                            "required": ["location"]
+                        }
+                    }
+                }
+            ]
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/parse/function_call")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+
+        // Parser endpoint should return 200 for valid requests
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Expected OK (200), got {}",
+            resp.status()
+        );
+
+        // Verify response contains tool_calls
+        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
+            .await
+            .unwrap();
+        let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
+        assert_eq!(body_json["success"], true);
+        assert!(body_json["tool_calls"].is_array());
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_parse_function_call_invalid_parser() {
+        let ctx = ParserTestContext::new(vec![]).await;
+        let app = ctx.create_app().await;
+
+        let payload = json!({
+            "text": "some text",
+            "tool_call_parser": "nonexistent_parser",
+            "tools": []
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/parse/function_call")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+
+        // Should return 400 (parser not found)
+        assert_eq!(
+            resp.status(),
+            StatusCode::BAD_REQUEST,
+            "Expected BAD_REQUEST (400), got {}",
+            resp.status()
+        );
+
+        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
+            .await
+            .unwrap();
+        let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
+
+        assert_eq!(body_json["success"], false);
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_parse_function_call_missing_fields() {
+        let ctx = ParserTestContext::new(vec![]).await;
+        let app = ctx.create_app().await;
+
+        // Missing 'text' field
+        let payload = json!({
+            "tool_call_parser": "json",
+            "tools": []
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/parse/function_call")
+            .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::UNPROCESSABLE_ENTITY);
+
+        // Missing 'tool_call_parser' field
+        let payload = json!({
+            "text": "some text",
+            "tools": []
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/parse/function_call")
+            .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::UNPROCESSABLE_ENTITY);
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_parse_function_call_empty_text() {
+        let ctx = ParserTestContext::new(vec![]).await;
+        let app = ctx.create_app().await;
+
+        let payload = json!({
+            "text": "",
+            "tool_call_parser": "json",
+            "tools": []
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/parse/function_call")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+
+        // Parser should handle empty text gracefully - return 200
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Expected OK (200), got {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod separate_reasoning_tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_separate_reasoning_success() {
+        let ctx = ParserTestContext::new(vec![]).await;
+        let app = ctx.create_app().await;
+
+        let payload = json!({
+            "text": "<think>Let me think about this problem. The user is asking for help.</think>Sure, I can help you with that.",
+            "reasoning_parser": "step3"
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/parse/reasoning")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+
+        // Should return 200 with parser factory initialized
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Expected OK (200), got {}",
+            resp.status()
+        );
+
+        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
+            .await
+            .unwrap();
+        let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
+
+        // Check response structure
+        assert_eq!(body_json["success"], true);
+        assert!(body_json.get("normal_text").is_some());
+        assert!(body_json.get("reasoning_text").is_some());
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_separate_reasoning_invalid_parser() {
+        let ctx = ParserTestContext::new(vec![]).await;
+        let app = ctx.create_app().await;
+
+        let payload = json!({
+            "text": "some text",
+            "reasoning_parser": "invalid_parser_type"
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/parse/reasoning")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+
+        // Should return 400 (parser not found)
+        assert_eq!(
+            resp.status(),
+            StatusCode::BAD_REQUEST,
+            "Expected BAD_REQUEST (400), got {}",
+            resp.status()
+        );
+
+        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
+            .await
+            .unwrap();
+        let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
+
+        assert_eq!(body_json["success"], false);
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_separate_reasoning_missing_fields() {
+        let ctx = ParserTestContext::new(vec![]).await;
+        let app = ctx.create_app().await;
+
+        // Missing 'text' field
+        let payload = json!({
+            "reasoning_parser": "step3"
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/parse/reasoning")
+            .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::UNPROCESSABLE_ENTITY);
+
+        // Missing 'reasoning_parser' field
+        let payload = json!({
+            "text": "some text"
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/parse/reasoning")
+            .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::UNPROCESSABLE_ENTITY);
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_separate_reasoning_empty_text() {
+        let ctx = ParserTestContext::new(vec![]).await;
+        let app = ctx.create_app().await;
+
+        let payload = json!({
+            "text": "",
+            "reasoning_parser": "step3"
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/parse/reasoning")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+
+        // Parser should handle empty text gracefully
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Expected OK (200), got {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_separate_reasoning_without_reasoning_tags() {
+        let ctx = ParserTestContext::new(vec![]).await;
+        let app = ctx.create_app().await;
+
+        let payload = json!({
+            "text": "Just a normal text without any reasoning tags",
+            "reasoning_parser": "step3"
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/parse/reasoning")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+
+        // Should return 200, parser should handle gracefully
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Expected OK (200), got {}",
+            resp.status()
+        );
+
+        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
+            .await
+            .unwrap();
+        let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
+
+        assert_eq!(body_json["success"], true);
+        // When there are no reasoning tags, parser returns empty normal_text and empty reasoning_text
+        // since the detect_and_parse_reasoning method only extracts if it finds reasoning markers
+        assert!(body_json.get("normal_text").is_some());
+        assert!(body_json.get("reasoning_text").is_some());
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_separate_reasoning_multiple_reasoning_blocks() {
+        let ctx = ParserTestContext::new(vec![]).await;
+        let app = ctx.create_app().await;
+
+        // Some parsers may handle multiple reasoning blocks
+        let payload = json!({
+            "text": "<think>First thought</think>Text 1<think>Second thought</think>Text 2",
+            "reasoning_parser": "step3"
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/parse/reasoning")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+
+        // Should handle multiple blocks gracefully
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Expected OK (200), got {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod api_routing_tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_admin_routes_accessible() {
+        let ctx = ParserTestContext::new(vec![]).await;
+        let app = ctx.create_app().await;
+
+        // Test that both endpoints exist and are accessible (even if parser factory not initialized)
+        let payload = json!({
+            "text": "test",
+            "tool_call_parser": "json",
+            "tools": []
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/parse/function_call")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.clone().oneshot(req).await.unwrap();
+
+        // Should not be 404
+        assert_ne!(resp.status(), StatusCode::NOT_FOUND);
+
+        let payload = json!({
+            "text": "test",
+            "reasoning_parser": "step3"
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/parse/reasoning")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+
+        // Should not be 404
+        assert_ne!(resp.status(), StatusCode::NOT_FOUND);
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_endpoints_only_accept_post() {
+        let ctx = ParserTestContext::new(vec![]).await;
+        let app = ctx.create_app().await;
+
+        // Test GET request to parse/function_call
+        let req = Request::builder()
+            .method("GET")
+            .uri("/parse/function_call")
+            .body(Body::empty())
+            .unwrap();
+
+        let resp = app.clone().oneshot(req).await.unwrap();
+
+        // Should not accept GET (should be 405 or 404)
+        assert!(
+            resp.status() == StatusCode::METHOD_NOT_ALLOWED
+                || resp.status() == StatusCode::NOT_FOUND,
+            "Expected METHOD_NOT_ALLOWED (405) or NOT_FOUND (404), got {}",
+            resp.status()
+        );
+
+        // Test GET request to parse/reasoning
+        let req = Request::builder()
+            .method("GET")
+            .uri("/parse/reasoning")
+            .body(Body::empty())
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+
+        // Should not accept GET
+        assert!(
+            resp.status() == StatusCode::METHOD_NOT_ALLOWED
+                || resp.status() == StatusCode::NOT_FOUND,
+            "Expected METHOD_NOT_ALLOWED (405) or NOT_FOUND (404), got {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+}

third_party/sglang/sgl-model-gateway/tests/api/request_formats_test.rs

@@ -0,0 +1,260 @@
+use serde_json::json;
+
+use crate::common::{
+    mock_worker::{HealthStatus, MockWorkerConfig, WorkerType},
+    WorkerTestContext,
+};
+
+#[cfg(test)]
+mod request_format_tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_generate_request_formats() {
+        let ctx = WorkerTestContext::new(vec![MockWorkerConfig {
+            port: 19001,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let payload = json!({
+            "text": "Hello, world!",
+            "stream": false
+        });
+
+        let result = ctx.make_request("/generate", payload).await;
+        assert!(result.is_ok());
+
+        let payload = json!({
+            "text": "Tell me a story",
+            "sampling_params": {
+                "temperature": 0.7,
+                "max_new_tokens": 100,
+                "top_p": 0.9
+            },
+            "stream": false
+        });
+
+        let result = ctx.make_request("/generate", payload).await;
+        assert!(result.is_ok());
+
+        let payload = json!({
+            "input_ids": [1, 2, 3, 4, 5],
+            "sampling_params": {
+                "temperature": 0.0,
+                "max_new_tokens": 50
+            },
+            "stream": false
+        });
+
+        let result = ctx.make_request("/generate", payload).await;
+        assert!(result.is_ok());
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_v1_chat_completions_formats() {
+        let ctx = WorkerTestContext::new(vec![MockWorkerConfig {
+            port: 19002,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let payload = json!({
+            "model": "test-model",
+            "messages": [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Hello!"}
+            ],
+            "stream": false
+        });
+
+        let result = ctx.make_request("/v1/chat/completions", payload).await;
+        assert!(result.is_ok());
+
+        let response = result.unwrap();
+        assert!(response.get("choices").is_some());
+        assert!(response.get("id").is_some());
+        assert_eq!(
+            response.get("object").and_then(|v| v.as_str()),
+            Some("chat.completion")
+        );
+
+        let payload = json!({
+            "model": "test-model",
+            "messages": [
+                {"role": "user", "content": "Tell me a joke"}
+            ],
+            "temperature": 0.8,
+            "max_tokens": 150,
+            "top_p": 0.95,
+            "stream": false
+        });
+
+        let result = ctx.make_request("/v1/chat/completions", payload).await;
+        assert!(result.is_ok());
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_v1_completions_formats() {
+        let ctx = WorkerTestContext::new(vec![MockWorkerConfig {
+            port: 19003,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let payload = json!({
+            "model": "test-model",
+            "prompt": "Once upon a time",
+            "max_tokens": 50,
+            "stream": false
+        });
+
+        let result = ctx.make_request("/v1/completions", payload).await;
+        assert!(result.is_ok());
+
+        let response = result.unwrap();
+        assert!(response.get("choices").is_some());
+        assert_eq!(
+            response.get("object").and_then(|v| v.as_str()),
+            Some("text_completion")
+        );
+
+        let payload = json!({
+            "model": "test-model",
+            "prompt": ["First prompt", "Second prompt"],
+            "temperature": 0.5,
+            "stream": false
+        });
+
+        let result = ctx.make_request("/v1/completions", payload).await;
+        assert!(result.is_ok());
+
+        let payload = json!({
+            "model": "test-model",
+            "prompt": "The capital of France is",
+            "max_tokens": 10,
+            "logprobs": 5,
+            "stream": false
+        });
+
+        let result = ctx.make_request("/v1/completions", payload).await;
+        assert!(result.is_ok());
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_batch_requests() {
+        let ctx = WorkerTestContext::new(vec![MockWorkerConfig {
+            port: 19004,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let payload = json!({
+            "text": ["First text", "Second text", "Third text"],
+            "sampling_params": {
+                "temperature": 0.7,
+                "max_new_tokens": 50
+            },
+            "stream": false
+        });
+
+        let result = ctx.make_request("/generate", payload).await;
+        assert!(result.is_ok());
+
+        let payload = json!({
+            "input_ids": [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+            "stream": false
+        });
+
+        let result = ctx.make_request("/generate", payload).await;
+        assert!(result.is_ok());
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_special_parameters() {
+        let ctx = WorkerTestContext::new(vec![MockWorkerConfig {
+            port: 19005,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let payload = json!({
+            "text": "Test",
+            "return_logprob": true,
+            "stream": false
+        });
+
+        let result = ctx.make_request("/generate", payload).await;
+        assert!(result.is_ok());
+
+        let payload = json!({
+            "text": "Generate JSON",
+            "sampling_params": {
+                "temperature": 0.0,
+                "json_schema": "$$ANY$$"
+            },
+            "stream": false
+        });
+
+        let result = ctx.make_request("/generate", payload).await;
+        assert!(result.is_ok());
+
+        let payload = json!({
+            "text": "Continue forever",
+            "sampling_params": {
+                "temperature": 0.7,
+                "max_new_tokens": 100,
+                "ignore_eos": true
+            },
+            "stream": false
+        });
+
+        let result = ctx.make_request("/generate", payload).await;
+        assert!(result.is_ok());
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_error_handling() {
+        let ctx = WorkerTestContext::new(vec![MockWorkerConfig {
+            port: 19006,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let payload = json!({});
+
+        let result = ctx.make_request("/generate", payload).await;
+        // Mock worker accepts empty body
+        assert!(result.is_ok());
+
+        ctx.shutdown().await;
+    }
+}

third_party/sglang/sgl-model-gateway/tests/api/streaming_tests.rs

@@ -0,0 +1,220 @@
+use serde_json::json;
+
+use crate::common::{
+    mock_worker::{HealthStatus, MockWorkerConfig, WorkerType},
+    WorkerTestContext,
+};
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_generate_streaming() {
+        let ctx = WorkerTestContext::new(vec![MockWorkerConfig {
+            port: 20001,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 10,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let payload = json!({
+            "text": "Stream test",
+            "stream": true,
+            "sampling_params": {
+                "temperature": 0.7,
+                "max_new_tokens": 10
+            }
+        });
+
+        let result = ctx.make_streaming_request("/generate", payload).await;
+        assert!(result.is_ok());
+
+        let events = result.unwrap();
+        assert!(events.len() >= 2);
+        assert_eq!(events.last().unwrap(), "[DONE]");
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_v1_chat_completions_streaming() {
+        let ctx = WorkerTestContext::new(vec![MockWorkerConfig {
+            port: 20002,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 10,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let payload = json!({
+            "model": "test-model",
+            "messages": [
+                {"role": "user", "content": "Count to 3"}
+            ],
+            "stream": true,
+            "max_tokens": 20
+        });
+
+        let result = ctx
+            .make_streaming_request("/v1/chat/completions", payload)
+            .await;
+        assert!(result.is_ok());
+
+        let events = result.unwrap();
+        assert!(events.len() >= 2);
+
+        for event in &events {
+            if event != "[DONE]" {
+                let parsed: Result<serde_json::Value, _> = serde_json::from_str(event);
+                assert!(parsed.is_ok(), "Invalid JSON in SSE event: {}", event);
+
+                let json = parsed.unwrap();
+                assert_eq!(
+                    json.get("object").and_then(|v| v.as_str()),
+                    Some("chat.completion.chunk")
+                );
+            }
+        }
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_v1_completions_streaming() {
+        let ctx = WorkerTestContext::new(vec![MockWorkerConfig {
+            port: 20003,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 10,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let payload = json!({
+            "model": "test-model",
+            "prompt": "Once upon a time",
+            "stream": true,
+            "max_tokens": 15
+        });
+
+        let result = ctx.make_streaming_request("/v1/completions", payload).await;
+        assert!(result.is_ok());
+
+        let events = result.unwrap();
+        assert!(events.len() >= 2);
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_streaming_with_error() {
+        let ctx = WorkerTestContext::new(vec![MockWorkerConfig {
+            port: 20004,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 1.0,
+        }])
+        .await;
+
+        let payload = json!({
+            "text": "This should fail",
+            "stream": true
+        });
+
+        let result = ctx.make_streaming_request("/generate", payload).await;
+        assert!(result.is_err());
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_streaming_timeouts() {
+        let ctx = WorkerTestContext::new(vec![MockWorkerConfig {
+            port: 20005,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 100,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let payload = json!({
+            "text": "Slow stream",
+            "stream": true,
+            "sampling_params": {
+                "max_new_tokens": 5
+            }
+        });
+
+        let start = std::time::Instant::now();
+        let result = ctx.make_streaming_request("/generate", payload).await;
+        let elapsed = start.elapsed();
+
+        assert!(result.is_ok());
+        let events = result.unwrap();
+        assert!(!events.is_empty());
+        assert!(elapsed.as_millis() >= 100);
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_batch_streaming() {
+        let ctx = WorkerTestContext::new(vec![MockWorkerConfig {
+            port: 20006,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 10,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let payload = json!({
+            "text": ["First", "Second", "Third"],
+            "stream": true,
+            "sampling_params": {
+                "max_new_tokens": 5
+            }
+        });
+
+        let result = ctx.make_streaming_request("/generate", payload).await;
+        assert!(result.is_ok());
+
+        let events = result.unwrap();
+        assert!(events.len() >= 4);
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_sse_format_parsing() {
+        let parse_sse_chunk = |chunk: &[u8]| -> Vec<String> {
+            let text = String::from_utf8_lossy(chunk);
+            text.lines()
+                .filter(|line| line.starts_with("data: "))
+                .map(|line| line[6..].to_string())
+                .collect()
+        };
+
+        let sse_data =
+            b"data: {\"text\":\"Hello\"}\n\ndata: {\"text\":\" world\"}\n\ndata: [DONE]\n\n";
+        let events = parse_sse_chunk(sse_data);
+
+        assert_eq!(events.len(), 3);
+        assert_eq!(events[0], "{\"text\":\"Hello\"}");
+        assert_eq!(events[1], "{\"text\":\" world\"}");
+        assert_eq!(events[2], "[DONE]");
+
+        let mixed = b"event: message\ndata: {\"test\":true}\n\n: comment\ndata: [DONE]\n\n";
+        let events = parse_sse_chunk(mixed);
+
+        assert_eq!(events.len(), 2);
+        assert_eq!(events[0], "{\"test\":true}");
+        assert_eq!(events[1], "[DONE]");
+    }
+}

third_party/sglang/sgl-model-gateway/tests/api_tests.rs

@@ -0,0 +1,6 @@
+//! API integration tests
+
+#[path = "common/mod.rs"]
+pub mod common;
+
+mod api;

third_party/sglang/sgl-model-gateway/tests/common/mock_mcp_server.rs

@@ -0,0 +1,180 @@
+// tests/common/mock_mcp_server.rs - Mock MCP server for testing
+use rmcp::{
+    handler::server::{router::tool::ToolRouter, wrapper::Parameters},
+    model::*,
+    service::RequestContext,
+    tool, tool_handler, tool_router,
+    transport::streamable_http_server::{
+        session::local::LocalSessionManager, StreamableHttpService,
+    },
+    ErrorData as McpError, RoleServer, ServerHandler,
+};
+use tokio::net::TcpListener;
+
+/// Mock MCP server that returns hardcoded responses for testing
+pub struct MockMCPServer {
+    pub port: u16,
+    pub server_handle: Option<tokio::task::JoinHandle<()>>,
+}
+
+/// Simple test server with mock search tools
+#[derive(Clone)]
+pub struct MockSearchServer {
+    tool_router: ToolRouter<MockSearchServer>,
+}
+
+impl Default for MockSearchServer {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[tool_router]
+impl MockSearchServer {
+    pub fn new() -> Self {
+        Self {
+            tool_router: Self::tool_router(),
+        }
+    }
+
+    #[tool(description = "Mock web search tool")]
+    fn brave_web_search(
+        &self,
+        Parameters(params): Parameters<serde_json::Map<String, serde_json::Value>>,
+    ) -> Result<CallToolResult, McpError> {
+        let query = params
+            .get("query")
+            .and_then(|v| v.as_str())
+            .unwrap_or("test");
+        Ok(CallToolResult::success(vec![Content::text(format!(
+            "Mock search results for: {}",
+            query
+        ))]))
+    }
+
+    #[tool(description = "Mock local search tool")]
+    fn brave_local_search(
+        &self,
+        Parameters(_params): Parameters<serde_json::Map<String, serde_json::Value>>,
+    ) -> Result<CallToolResult, McpError> {
+        Ok(CallToolResult::success(vec![Content::text(
+            "Mock local search results",
+        )]))
+    }
+}
+
+#[tool_handler]
+impl ServerHandler for MockSearchServer {
+    fn get_info(&self) -> ServerInfo {
+        ServerInfo {
+            protocol_version: ProtocolVersion::V_2024_11_05,
+            capabilities: ServerCapabilities::builder().enable_tools().build(),
+            server_info: Implementation::from_build_env(),
+            instructions: Some("Mock server for testing".to_string()),
+        }
+    }
+
+    async fn initialize(
+        &self,
+        _request: InitializeRequestParam,
+        _context: RequestContext<RoleServer>,
+    ) -> Result<InitializeResult, McpError> {
+        Ok(self.get_info())
+    }
+}
+
+impl MockMCPServer {
+    /// Start a mock MCP server on an available port
+    pub async fn start() -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
+        // Find an available port
+        let listener = TcpListener::bind("127.0.0.1:0").await?;
+        let port = listener.local_addr()?.port();
+
+        // Create the MCP service using rmcp's StreamableHttpService
+        let service = StreamableHttpService::new(
+            || Ok(MockSearchServer::new()),
+            LocalSessionManager::default().into(),
+            Default::default(),
+        );
+
+        let app = axum::Router::new().nest_service("/mcp", service);
+
+        let server_handle = tokio::spawn(async move {
+            axum::serve(listener, app)
+                .await
+                .expect("Mock MCP server failed to start");
+        });
+
+        // Give the server a moment to start
+        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+
+        Ok(MockMCPServer {
+            port,
+            server_handle: Some(server_handle),
+        })
+    }
+
+    /// Get the full URL for this mock server
+    pub fn url(&self) -> String {
+        format!("http://127.0.0.1:{}/mcp", self.port)
+    }
+
+    /// Stop the mock server
+    pub async fn stop(&mut self) {
+        if let Some(handle) = self.server_handle.take() {
+            handle.abort();
+            // Wait a moment for cleanup
+            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+        }
+    }
+}
+
+impl Drop for MockMCPServer {
+    fn drop(&mut self) {
+        if let Some(handle) = self.server_handle.take() {
+            handle.abort();
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[allow(unused_imports)]
+    use super::MockMCPServer;
+
+    #[tokio::test]
+    async fn test_mock_server_startup() {
+        let mut server = MockMCPServer::start().await.unwrap();
+        assert!(server.port > 0);
+        assert!(server.url().contains(&server.port.to_string()));
+        server.stop().await;
+    }
+
+    #[tokio::test]
+    async fn test_mock_server_with_rmcp_client() {
+        let mut server = MockMCPServer::start().await.unwrap();
+
+        use rmcp::{transport::StreamableHttpClientTransport, ServiceExt};
+
+        let transport = StreamableHttpClientTransport::from_uri(server.url().as_str());
+        let client = ().serve(transport).await;
+
+        assert!(client.is_ok(), "Should be able to connect to mock server");
+
+        if let Ok(client) = client {
+            let tools = client.peer().list_all_tools().await;
+            assert!(tools.is_ok(), "Should be able to list tools");
+
+            if let Ok(tools) = tools {
+                assert_eq!(tools.len(), 2, "Should have 2 tools");
+                assert!(tools.iter().any(|t| t.name == "brave_web_search"));
+                assert!(tools.iter().any(|t| t.name == "brave_local_search"));
+            }
+
+            // Shutdown by dropping the client
+            drop(client);
+        }
+
+        server.stop().await;
+    }
+}

third_party/sglang/sgl-model-gateway/tests/common/mock_openai_server.rs

@@ -0,0 +1,240 @@
+//! Mock servers for testing
+
+#![allow(dead_code)]
+
+use std::{net::SocketAddr, sync::Arc};
+
+use axum::{
+    body::Body,
+    extract::{Request, State},
+    http::{HeaderValue, StatusCode},
+    response::{
+        sse::{Event, KeepAlive},
+        IntoResponse, Response, Sse,
+    },
+    routing::post,
+    Json, Router,
+};
+use futures_util::stream::{self, StreamExt};
+use serde_json::json;
+use tokio::net::TcpListener;
+
+/// Mock OpenAI API server for testing
+pub struct MockOpenAIServer {
+    addr: SocketAddr,
+    _handle: tokio::task::JoinHandle<()>,
+}
+
+#[derive(Clone)]
+struct MockServerState {
+    require_auth: bool,
+    expected_auth: Option<String>,
+}
+
+impl MockOpenAIServer {
+    /// Create and start a new mock OpenAI server
+    pub async fn new() -> Self {
+        Self::new_with_auth(None).await
+    }
+
+    /// Create and start a new mock OpenAI server with optional auth requirement
+    pub async fn new_with_auth(expected_auth: Option<String>) -> Self {
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let addr = listener.local_addr().unwrap();
+
+        let state = Arc::new(MockServerState {
+            require_auth: expected_auth.is_some(),
+            expected_auth,
+        });
+
+        let app = Router::new()
+            .route("/v1/chat/completions", post(mock_chat_completions))
+            .route("/v1/completions", post(mock_completions))
+            .route("/v1/models", post(mock_models).get(mock_models))
+            .with_state(state);
+
+        let handle = tokio::spawn(async move {
+            axum::serve(listener, app).await.unwrap();
+        });
+
+        // Give the server a moment to start
+        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
+
+        Self {
+            addr,
+            _handle: handle,
+        }
+    }
+
+    /// Get the base URL for this mock server
+    pub fn base_url(&self) -> String {
+        format!("http://{}", self.addr)
+    }
+}
+
+/// Mock chat completions endpoint
+async fn mock_chat_completions(req: Request<Body>) -> Response {
+    let (_, body) = req.into_parts();
+    let body_bytes = match axum::body::to_bytes(body, usize::MAX).await {
+        Ok(bytes) => bytes,
+        Err(_) => return StatusCode::BAD_REQUEST.into_response(),
+    };
+
+    let request: serde_json::Value = match serde_json::from_slice(&body_bytes) {
+        Ok(req) => req,
+        Err(_) => return StatusCode::BAD_REQUEST.into_response(),
+    };
+
+    // Extract model from request or use default (owned String to satisfy 'static in stream)
+    let model: String = request
+        .get("model")
+        .and_then(|v| v.as_str())
+        .unwrap_or("gpt-3.5-turbo")
+        .to_string();
+
+    // If stream requested, return SSE
+    let is_stream = request
+        .get("stream")
+        .and_then(|v| v.as_bool())
+        .unwrap_or(false);
+
+    if is_stream {
+        let created = 1677652288u64;
+        // Single chunk then [DONE]
+        let model_chunk = model.clone();
+        let event_stream = stream::once(async move {
+            let chunk = json!({
+                "id": "chatcmpl-123456789",
+                "object": "chat.completion.chunk",
+                "created": created,
+                "model": model_chunk,
+                "choices": [{
+                    "index": 0,
+                    "delta": {
+                        "content": "Hello!"
+                    },
+                    "finish_reason": null
+                }]
+            });
+            Ok::<_, std::convert::Infallible>(Event::default().data(chunk.to_string()))
+        })
+        .chain(stream::once(async { Ok(Event::default().data("[DONE]")) }));
+
+        Sse::new(event_stream)
+            .keep_alive(KeepAlive::default())
+            .into_response()
+    } else {
+        // Create a mock non-streaming response
+        let response = json!({
+            "id": "chatcmpl-123456789",
+            "object": "chat.completion",
+            "created": 1677652288,
+            "model": model,
+            "choices": [{
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": "Hello! I'm a mock OpenAI assistant. How can I help you today?"
+                },
+                "finish_reason": "stop"
+            }],
+            "usage": {
+                "prompt_tokens": 9,
+                "completion_tokens": 12,
+                "total_tokens": 21
+            }
+        });
+
+        Json(response).into_response()
+    }
+}
+
+/// Mock completions endpoint (legacy)
+async fn mock_completions(req: Request<Body>) -> Response {
+    let (_, body) = req.into_parts();
+    let body_bytes = match axum::body::to_bytes(body, usize::MAX).await {
+        Ok(bytes) => bytes,
+        Err(_) => return StatusCode::BAD_REQUEST.into_response(),
+    };
+
+    let request: serde_json::Value = match serde_json::from_slice(&body_bytes) {
+        Ok(req) => req,
+        Err(_) => return StatusCode::BAD_REQUEST.into_response(),
+    };
+
+    let model = request["model"].as_str().unwrap_or("text-davinci-003");
+
+    let response = json!({
+        "id": "cmpl-123456789",
+        "object": "text_completion",
+        "created": 1677652288,
+        "model": model,
+        "choices": [{
+            "text": " This is a mock completion response.",
+            "index": 0,
+            "logprobs": null,
+            "finish_reason": "stop"
+        }],
+        "usage": {
+            "prompt_tokens": 5,
+            "completion_tokens": 7,
+            "total_tokens": 12
+        }
+    });
+
+    Json(response).into_response()
+}
+
+/// Mock models endpoint
+async fn mock_models(State(state): State<Arc<MockServerState>>, req: Request<Body>) -> Response {
+    // Optionally enforce Authorization header
+    if state.require_auth {
+        let auth = req
+            .headers()
+            .get("authorization")
+            .or_else(|| req.headers().get("Authorization"))
+            .and_then(|v| v.to_str().ok())
+            .map(|s| s.to_string());
+        let auth_ok = match (&state.expected_auth, auth) {
+            (Some(expected), Some(got)) => &got == expected,
+            (None, Some(_)) => true,
+            _ => false,
+        };
+        if !auth_ok {
+            let mut response = Response::new(Body::from(
+                json!({
+                    "error": {
+                        "message": "Unauthorized",
+                        "type": "invalid_request_error"
+                    }
+                })
+                .to_string(),
+            ));
+            *response.status_mut() = StatusCode::UNAUTHORIZED;
+            response
+                .headers_mut()
+                .insert("WWW-Authenticate", HeaderValue::from_static("Bearer"));
+            return response;
+        }
+    }
+
+    let response = json!({
+        "object": "list",
+        "data": [
+            {
+                "id": "gpt-4",
+                "object": "model",
+                "created": 1677610602,
+                "owned_by": "openai"
+            },
+            {
+                "id": "gpt-3.5-turbo",
+                "object": "model",
+                "created": 1677610602,
+                "owned_by": "openai"
+            }
+        ]
+    });
+
+    Json(response).into_response()
+}

third_party/sglang/sgl-model-gateway/tests/common/redis_test_server.rs

@@ -0,0 +1,118 @@
+use std::{
+    process::{Child, Command},
+    sync::OnceLock,
+    time::Duration,
+};
+
+use redis::RedisError;
+use tracing::{info, warn};
+
+static SHARED_SERVER: OnceLock<RedisTestServer> = OnceLock::new();
+
+pub fn get_shared_server() -> &'static RedisTestServer {
+    let server = SHARED_SERVER
+        .get_or_init(|| RedisTestServer::start().expect("Failed to start shared Redis server"));
+    server.wait_ready();
+    server
+}
+
+pub struct RedisTestServer {
+    process: Option<Child>,
+    port: u16,
+    url: String,
+}
+
+impl RedisTestServer {
+    pub fn start() -> Result<Self, String> {
+        let port = portpicker::pick_unused_port()
+            .ok_or_else(|| "Failed to find available port".to_string())?;
+        let url = format!("redis://127.0.0.1:{}", port);
+
+        let mut cmd = Command::new("redis-server");
+        cmd.args([
+            "--port",
+            &port.to_string(),
+            "--save",
+            "",
+            "--appendonly",
+            "no",
+            "--daemonize",
+            "no",
+        ]);
+
+        info!("Starting redis server... cmd={cmd:?}");
+        let process = cmd.spawn().map_err(|e| {
+            format!(
+                "Failed to start redis-server: {}. Is redis-server installed?",
+                e
+            )
+        })?;
+
+        Ok(Self {
+            process: Some(process),
+            port,
+            url,
+        })
+    }
+
+    pub fn wait_ready(&self) {
+        for _ in 0..200 {
+            match self.is_ready() {
+                Ok(()) => return,
+                Err(e) => info!("wait_ready failed, will retry (e={e})"),
+            }
+            std::thread::sleep(Duration::from_millis(100));
+        }
+        panic!("Timeout waiting Redis server ready on port {}", self.port);
+    }
+
+    pub fn is_ready(&self) -> Result<(), RedisError> {
+        let client = redis::Client::open(self.url.as_str())?;
+        let mut conn = client.get_connection()?;
+        redis::cmd("PING").query::<String>(&mut conn)?;
+        info!("Redis is_ready=true url={}", self.url);
+        Ok(())
+    }
+
+    pub fn url(&self) -> &str {
+        &self.url
+    }
+}
+
+impl Drop for RedisTestServer {
+    fn drop(&mut self) {
+        if let Some(mut process) = self.process.take() {
+            info!("Killing redis server... process={process:?}");
+            if let Err(e) = process.kill() {
+                warn!("Failed to kill redis-server process: {}", e);
+            }
+            if let Err(e) = process.wait() {
+                warn!("Failed to wait for redis-server process: {}", e);
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_redis_server_start_stop() {
+        let server = RedisTestServer::start().unwrap();
+        server.wait_ready();
+        assert!(server.url().starts_with("redis://"));
+
+        let client = redis::Client::open(server.url()).unwrap();
+        let mut conn = client.get_connection().unwrap();
+
+        let _: () = redis::cmd("SET")
+            .arg("test_key")
+            .arg("test_value")
+            .query(&mut conn)
+            .unwrap();
+
+        let value: String = redis::cmd("GET").arg("test_key").query(&mut conn).unwrap();
+        assert_eq!(value, "test_value");
+    }
+}

third_party/sglang/sgl-model-gateway/tests/common/streaming_helpers.rs

@@ -0,0 +1,134 @@
+//! Streaming Test Helpers
+//!
+//! Utilities for creating realistic streaming chunks that simulate
+//! how LLM tokens actually arrive (1-5 characters at a time).
+
+/// Split input into realistic char-level chunks (2-3 chars each for determinism)
+pub fn create_realistic_chunks(input: &str) -> Vec<String> {
+    let mut chunks = Vec::new();
+    let chars: Vec<char> = input.chars().collect();
+    let mut i = 0;
+
+    while i < chars.len() {
+        // Take 2-3 characters at a time (deterministic for testing)
+        let chunk_size = if i + 3 <= chars.len() && chars[i].is_ascii_alphanumeric() {
+            3 // Longer chunks for alphanumeric sequences
+        } else {
+            2 // Shorter chunks for special characters
+        };
+
+        let end = (i + chunk_size).min(chars.len());
+        let chunk: String = chars[i..end].iter().collect();
+        chunks.push(chunk);
+        i = end;
+    }
+
+    chunks
+}
+
+/// Split input at strategic positions to test edge cases
+/// This creates chunks that break at critical positions like after quotes, colons, etc.
+pub fn create_strategic_chunks(input: &str) -> Vec<String> {
+    let mut chunks = Vec::new();
+    let mut current = String::new();
+    let chars: Vec<char> = input.chars().collect();
+
+    for (i, &ch) in chars.iter().enumerate() {
+        current.push(ch);
+
+        // Break after strategic characters
+        let should_break = matches!(ch, '"' | ':' | ',' | '{' | '}' | '[' | ']')
+            || (i > 0 && chars[i-1] == '"' && ch == ' ') // Space after quote
+            || current.len() >= 5; // Max 5 chars per chunk
+
+        if should_break && !current.is_empty() {
+            chunks.push(current.clone());
+            current.clear();
+        }
+    }
+
+    if !current.is_empty() {
+        chunks.push(current);
+    }
+
+    chunks
+}
+
+/// Create the bug scenario chunks: `{"name": "` arrives in parts
+pub fn create_bug_scenario_chunks() -> Vec<&'static str> {
+    vec![
+        r#"{"#,
+        r#"""#,
+        r#"name"#,
+        r#"""#,
+        r#":"#,
+        r#" "#,
+        r#"""#,      // Bug occurs here: parser has {"name": "
+        r#"search"#, // Use valid tool name
+        r#"""#,
+        r#","#,
+        r#" "#,
+        r#"""#,
+        r#"arguments"#,
+        r#"""#,
+        r#":"#,
+        r#" "#,
+        r#"{"#,
+        r#"""#,
+        r#"query"#,
+        r#"""#,
+        r#":"#,
+        r#" "#,
+        r#"""#,
+        r#"test query"#,
+        r#"""#,
+        r#"}"#,
+        r#"}"#,
+    ]
+}
+
+#[cfg(test)]
+mod tests {
+    #[allow(unused_imports)]
+    use super::*;
+
+    #[test]
+    fn test_realistic_chunks() {
+        let input = r#"{"name": "test"}"#;
+        let chunks = create_realistic_chunks(input);
+
+        // Should have multiple chunks
+        assert!(chunks.len() > 3);
+
+        // Reconstructed should equal original
+        let reconstructed: String = chunks.join("");
+        assert_eq!(reconstructed, input);
+    }
+
+    #[test]
+    fn test_strategic_chunks_breaks_after_quotes() {
+        let input = r#"{"name": "value"}"#;
+        let chunks = create_strategic_chunks(input);
+
+        // Should break after quotes and colons
+        assert!(chunks.iter().any(|c| c.ends_with('"')));
+        assert!(chunks.iter().any(|c| c.ends_with(':')));
+
+        // Reconstructed should equal original
+        let reconstructed: String = chunks.join("");
+        assert_eq!(reconstructed, input);
+    }
+
+    #[test]
+    fn test_bug_scenario_chunks() {
+        let chunks = create_bug_scenario_chunks();
+        let reconstructed: String = chunks.join("");
+
+        // Should reconstruct to valid JSON
+        assert!(reconstructed.contains(r#"{"name": "search""#));
+
+        // The critical chunk sequence should be present (space after colon, then quote in next chunk)
+        let joined = chunks.join("|");
+        assert!(joined.contains(r#" |"#)); // The bug happens at {"name": " and then "
+    }
+}

third_party/sglang/sgl-model-gateway/tests/common/test_app.rs

@@ -0,0 +1,267 @@
+use std::sync::{Arc, OnceLock};
+
+use axum::Router;
+use data_connector::{
+    MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage,
+};
+use reqwest::Client;
+use smg::{
+    app_context::AppContext,
+    config::RouterConfig,
+    core::{
+        BasicWorkerBuilder, LoadMonitor, ModelCard, RuntimeType, Worker, WorkerRegistry, WorkerType,
+    },
+    middleware::{AuthConfig, TokenBucket},
+    policies::PolicyRegistry,
+    routers::RouterTrait,
+    server::{build_app, AppState},
+    tokenizer::registry::TokenizerRegistry,
+};
+use smg_mcp::{McpConfig, McpManager};
+
+/// Create a test Axum application using the actual server's build_app function
+#[allow(dead_code)]
+pub fn create_test_app(
+    router: Arc<dyn RouterTrait>,
+    client: Client,
+    router_config: &RouterConfig,
+) -> Router {
+    // Initialize rate limiter
+    let rate_limiter = match router_config.max_concurrent_requests {
+        n if n <= 0 => None,
+        n => {
+            let rate_limit_tokens = router_config
+                .rate_limit_tokens_per_second
+                .filter(|&t| t > 0)
+                .unwrap_or(n);
+            Some(Arc::new(TokenBucket::new(
+                n as usize,
+                rate_limit_tokens as usize,
+            )))
+        }
+    };
+
+    // Initialize registries
+    let worker_registry = Arc::new(WorkerRegistry::new());
+    let policy_registry = Arc::new(PolicyRegistry::new(router_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(),
+        router_config.worker_startup_check_interval_secs,
+    )));
+
+    // Create empty OnceLock for worker job queue and workflow engines
+    let worker_job_queue = Arc::new(OnceLock::new());
+    let workflow_engines = Arc::new(OnceLock::new());
+
+    // Create AppContext using builder pattern
+    let app_context = Arc::new(
+        AppContext::builder()
+            .router_config(router_config.clone())
+            .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_engines(workflow_engines)
+            .build()
+            .unwrap(),
+    );
+
+    // Create AppState with the test router and context
+    let app_state = Arc::new(AppState {
+        router,
+        context: app_context,
+        concurrency_queue_tx: None,
+        router_manager: None,
+        mesh_handler: None,
+        mesh_sync_manager: None,
+    });
+
+    // Configure request ID headers (use defaults if not specified)
+    let request_id_headers = router_config.request_id_headers.clone().unwrap_or_else(|| {
+        vec![
+            "x-request-id".to_string(),
+            "x-correlation-id".to_string(),
+            "x-trace-id".to_string(),
+            "request-id".to_string(),
+        ]
+    });
+
+    // Create auth config from router config
+    let auth_config = AuthConfig {
+        api_key: router_config.api_key.clone(),
+    };
+
+    // Use the actual server's build_app function
+    build_app(
+        app_state,
+        auth_config,
+        None, // No control plane auth for tests
+        router_config.max_payload_size,
+        request_id_headers,
+        router_config.cors_allowed_origins.clone(),
+    )
+}
+
+/// Create a test Axum application with an existing AppContext
+#[allow(dead_code)]
+pub fn create_test_app_with_context(
+    router: Arc<dyn RouterTrait>,
+    app_context: Arc<AppContext>,
+) -> Router {
+    // Create AppState with the test router and context
+    let app_state = Arc::new(AppState {
+        router,
+        context: app_context.clone(),
+        concurrency_queue_tx: None,
+        router_manager: None,
+        mesh_handler: None,
+        mesh_sync_manager: None,
+    });
+
+    // Get config from the context
+    let router_config = &app_context.router_config;
+
+    // Configure request ID headers (use defaults if not specified)
+    let request_id_headers = router_config.request_id_headers.clone().unwrap_or_else(|| {
+        vec![
+            "x-request-id".to_string(),
+            "x-correlation-id".to_string(),
+            "x-trace-id".to_string(),
+            "request-id".to_string(),
+        ]
+    });
+
+    // Create auth config from router config
+    let auth_config = AuthConfig {
+        api_key: router_config.api_key.clone(),
+    };
+
+    // Use the actual server's build_app function
+    build_app(
+        app_state,
+        auth_config,
+        None, // No control plane auth for tests
+        router_config.max_payload_size,
+        request_id_headers,
+        router_config.cors_allowed_origins.clone(),
+    )
+}
+
+/// Create a minimal test AppContext for unit tests
+#[allow(dead_code)]
+pub async fn create_test_app_context() -> Arc<AppContext> {
+    let router_config = RouterConfig::default();
+    let client = Client::new();
+
+    // Initialize empty OnceLocks
+    let worker_job_queue = Arc::new(OnceLock::new());
+    let workflow_engines = Arc::new(OnceLock::new());
+
+    // Initialize MCP manager with empty config
+    let mcp_manager_lock = Arc::new(OnceLock::new());
+    let empty_config = McpConfig {
+        servers: vec![],
+        pool: Default::default(),
+        proxy: None,
+        warmup: vec![],
+        inventory: Default::default(),
+    };
+    let mcp_manager = McpManager::with_defaults(empty_config)
+        .await
+        .expect("Failed to create MCP manager");
+    mcp_manager_lock.set(Arc::new(mcp_manager)).ok();
+
+    // Initialize registries
+    let worker_registry = Arc::new(WorkerRegistry::new());
+    let policy_registry = Arc::new(PolicyRegistry::new(router_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());
+
+    Arc::new(
+        AppContext::builder()
+            .router_config(router_config)
+            .client(client)
+            .rate_limiter(None)
+            .tokenizer_registry(Arc::new(TokenizerRegistry::new()))
+            .reasoning_parser_factory(None)
+            .tool_parser_factory(None)
+            .worker_registry(worker_registry)
+            .policy_registry(policy_registry)
+            .response_storage(response_storage)
+            .conversation_storage(conversation_storage)
+            .conversation_item_storage(conversation_item_storage)
+            .load_monitor(None)
+            .worker_job_queue(worker_job_queue)
+            .workflow_engines(workflow_engines)
+            .mcp_manager(mcp_manager_lock)
+            .build()
+            .unwrap(),
+    )
+}
+
+/// Register an external worker (OpenAI-compatible API endpoint) in the test AppContext.
+///
+/// This is used by tests that need to test the OpenAI router, which expects
+/// workers to be registered in the WorkerRegistry before routing requests.
+///
+/// # Arguments
+/// * `ctx` - The AppContext to register the worker in
+/// * `url` - The base URL of the external API endpoint
+/// * `models` - Optional list of model IDs this worker supports. If empty, uses "gpt-3.5-turbo" as default.
+#[allow(dead_code)]
+pub fn register_external_worker(ctx: &Arc<AppContext>, url: &str, models: Option<Vec<&str>>) {
+    let model_list: Vec<ModelCard> = models
+        .unwrap_or_else(|| vec!["gpt-3.5-turbo"])
+        .into_iter()
+        .map(ModelCard::new)
+        .collect();
+
+    let worker: Arc<dyn Worker> = Arc::new(
+        BasicWorkerBuilder::new(url)
+            .worker_type(WorkerType::Regular)
+            .runtime_type(RuntimeType::External)
+            .models(model_list)
+            .build(),
+    );
+
+    ctx.worker_registry.register(worker);
+}
+
+/// Register an external worker with a custom model card that has aliases.
+///
+/// # Arguments
+/// * `ctx` - The AppContext to register the worker in
+/// * `url` - The base URL of the external API endpoint
+/// * `model_card` - A fully configured ModelCard with aliases, provider, etc.
+#[allow(dead_code)]
+pub fn register_external_worker_with_card(ctx: &Arc<AppContext>, url: &str, model_card: ModelCard) {
+    let worker: Arc<dyn Worker> = Arc::new(
+        BasicWorkerBuilder::new(url)
+            .worker_type(WorkerType::Regular)
+            .runtime_type(RuntimeType::External)
+            .model(model_card)
+            .build(),
+    );
+
+    ctx.worker_registry.register(worker);
+}

third_party/sglang/sgl-model-gateway/tests/common/test_certs.rs

@@ -0,0 +1,299 @@
+// TLS certificate generation for integration tests
+#![allow(dead_code)]
+
+use std::path::PathBuf;
+
+use openssl::{
+    asn1::Asn1Time,
+    bn::{BigNum, MsbOption},
+    hash::MessageDigest,
+    pkey::{PKey, Private},
+    rsa::Rsa,
+    x509::{
+        extension::{BasicConstraints, ExtendedKeyUsage, KeyUsage, SubjectAlternativeName},
+        X509Builder, X509NameBuilder, X509,
+    },
+};
+use tempfile::TempDir;
+
+/// Container for generated test certificates
+pub struct TestCertificates {
+    /// Temporary directory containing the certificate files
+    pub temp_dir: TempDir,
+    /// Path to CA certificate
+    pub ca_cert_path: PathBuf,
+    /// Path to CA private key
+    pub ca_key_path: PathBuf,
+    /// Path to server certificate
+    pub server_cert_path: PathBuf,
+    /// Path to server private key
+    pub server_key_path: PathBuf,
+    /// Path to client certificate
+    pub client_cert_path: PathBuf,
+    /// Path to client private key
+    pub client_key_path: PathBuf,
+}
+
+impl TestCertificates {
+    /// Generate a complete set of test certificates for mTLS testing.
+    ///
+    /// Creates:
+    /// - CA certificate and key (self-signed root CA)
+    /// - Server certificate and key (signed by CA, for localhost)
+    /// - Client certificate and key (signed by CA, for client authentication)
+    pub fn generate() -> Result<Self, Box<dyn std::error::Error>> {
+        let temp_dir = TempDir::new()?;
+        let base_path = temp_dir.path();
+
+        // Generate CA key pair
+        let ca_key = generate_rsa_key()?;
+        let ca_cert = generate_ca_certificate(&ca_key)?;
+
+        // Generate server key pair and certificate
+        let server_key = generate_rsa_key()?;
+        let server_cert = generate_server_certificate(&server_key, &ca_cert, &ca_key)?;
+
+        // Generate client key pair and certificate
+        let client_key = generate_rsa_key()?;
+        let client_cert = generate_client_certificate(&client_key, &ca_cert, &ca_key)?;
+
+        // Write all files
+        let ca_cert_path = base_path.join("ca_cert.pem");
+        let ca_key_path = base_path.join("ca_key.pem");
+        let server_cert_path = base_path.join("server_cert.pem");
+        let server_key_path = base_path.join("server_key.pem");
+        let client_cert_path = base_path.join("client_cert.pem");
+        let client_key_path = base_path.join("client_key.pem");
+
+        std::fs::write(&ca_cert_path, ca_cert.to_pem()?)?;
+        std::fs::write(&ca_key_path, ca_key.private_key_to_pem_pkcs8()?)?;
+        std::fs::write(&server_cert_path, server_cert.to_pem()?)?;
+        std::fs::write(&server_key_path, server_key.private_key_to_pem_pkcs8()?)?;
+        std::fs::write(&client_cert_path, client_cert.to_pem()?)?;
+        std::fs::write(&client_key_path, client_key.private_key_to_pem_pkcs8()?)?;
+
+        Ok(Self {
+            temp_dir,
+            ca_cert_path,
+            ca_key_path,
+            server_cert_path,
+            server_key_path,
+            client_cert_path,
+            client_key_path,
+        })
+    }
+
+    /// Get paths as string references for use with RouterConfig builder
+    pub fn ca_cert_str(&self) -> &str {
+        self.ca_cert_path.to_str().unwrap()
+    }
+
+    pub fn server_cert_str(&self) -> &str {
+        self.server_cert_path.to_str().unwrap()
+    }
+
+    pub fn server_key_str(&self) -> &str {
+        self.server_key_path.to_str().unwrap()
+    }
+
+    pub fn client_cert_str(&self) -> &str {
+        self.client_cert_path.to_str().unwrap()
+    }
+
+    pub fn client_key_str(&self) -> &str {
+        self.client_key_path.to_str().unwrap()
+    }
+}
+
+/// Generate a 2048-bit RSA key pair
+fn generate_rsa_key() -> Result<PKey<Private>, Box<dyn std::error::Error>> {
+    let rsa = Rsa::generate(2048)?;
+    Ok(PKey::from_rsa(rsa)?)
+}
+
+/// Generate a self-signed CA certificate
+fn generate_ca_certificate(key: &PKey<Private>) -> Result<X509, Box<dyn std::error::Error>> {
+    let mut name_builder = X509NameBuilder::new()?;
+    name_builder.append_entry_by_text("C", "US")?;
+    name_builder.append_entry_by_text("ST", "California")?;
+    name_builder.append_entry_by_text("L", "Test City")?;
+    name_builder.append_entry_by_text("O", "Test CA Organization")?;
+    name_builder.append_entry_by_text("CN", "Test CA")?;
+    let name = name_builder.build();
+
+    let mut cert_builder = X509Builder::new()?;
+    cert_builder.set_version(2)?; // X509 v3
+
+    // Serial number
+    let serial = {
+        let mut bn = BigNum::new()?;
+        bn.rand(128, MsbOption::MAYBE_ZERO, false)?;
+        bn.to_asn1_integer()?
+    };
+    cert_builder.set_serial_number(&serial)?;
+
+    cert_builder.set_subject_name(&name)?;
+    cert_builder.set_issuer_name(&name)?; // Self-signed
+    cert_builder.set_pubkey(key)?;
+
+    // Validity: 1 year from now
+    let not_before = Asn1Time::days_from_now(0)?;
+    let not_after = Asn1Time::days_from_now(365)?;
+    cert_builder.set_not_before(&not_before)?;
+    cert_builder.set_not_after(&not_after)?;
+
+    // Extensions for CA
+    let basic_constraints = BasicConstraints::new().critical().ca().build()?;
+    cert_builder.append_extension(basic_constraints)?;
+
+    let key_usage = KeyUsage::new()
+        .critical()
+        .key_cert_sign()
+        .crl_sign()
+        .build()?;
+    cert_builder.append_extension(key_usage)?;
+
+    cert_builder.sign(key, MessageDigest::sha256())?;
+
+    Ok(cert_builder.build())
+}
+
+/// Generate a server certificate signed by the CA
+fn generate_server_certificate(
+    key: &PKey<Private>,
+    ca_cert: &X509,
+    ca_key: &PKey<Private>,
+) -> Result<X509, Box<dyn std::error::Error>> {
+    let mut name_builder = X509NameBuilder::new()?;
+    name_builder.append_entry_by_text("C", "US")?;
+    name_builder.append_entry_by_text("ST", "California")?;
+    name_builder.append_entry_by_text("L", "Test City")?;
+    name_builder.append_entry_by_text("O", "Test Server Organization")?;
+    name_builder.append_entry_by_text("CN", "localhost")?;
+    let name = name_builder.build();
+
+    let mut cert_builder = X509Builder::new()?;
+    cert_builder.set_version(2)?;
+
+    let serial = {
+        let mut bn = BigNum::new()?;
+        bn.rand(128, MsbOption::MAYBE_ZERO, false)?;
+        bn.to_asn1_integer()?
+    };
+    cert_builder.set_serial_number(&serial)?;
+
+    cert_builder.set_subject_name(&name)?;
+    cert_builder.set_issuer_name(ca_cert.subject_name())?;
+    cert_builder.set_pubkey(key)?;
+
+    let not_before = Asn1Time::days_from_now(0)?;
+    let not_after = Asn1Time::days_from_now(365)?;
+    cert_builder.set_not_before(&not_before)?;
+    cert_builder.set_not_after(&not_after)?;
+
+    // Extensions for server certificate
+    let basic_constraints = BasicConstraints::new().build()?;
+    cert_builder.append_extension(basic_constraints)?;
+
+    let key_usage = KeyUsage::new()
+        .critical()
+        .digital_signature()
+        .key_encipherment()
+        .build()?;
+    cert_builder.append_extension(key_usage)?;
+
+    let ext_key_usage = ExtendedKeyUsage::new().server_auth().build()?;
+    cert_builder.append_extension(ext_key_usage)?;
+
+    // Subject Alternative Names for localhost
+    let san = SubjectAlternativeName::new()
+        .dns("localhost")
+        .ip("127.0.0.1")
+        .ip("::1")
+        .build(&cert_builder.x509v3_context(Some(ca_cert), None))?;
+    cert_builder.append_extension(san)?;
+
+    cert_builder.sign(ca_key, MessageDigest::sha256())?;
+
+    Ok(cert_builder.build())
+}
+
+/// Generate a client certificate signed by the CA
+fn generate_client_certificate(
+    key: &PKey<Private>,
+    ca_cert: &X509,
+    ca_key: &PKey<Private>,
+) -> Result<X509, Box<dyn std::error::Error>> {
+    let mut name_builder = X509NameBuilder::new()?;
+    name_builder.append_entry_by_text("C", "US")?;
+    name_builder.append_entry_by_text("ST", "California")?;
+    name_builder.append_entry_by_text("L", "Test City")?;
+    name_builder.append_entry_by_text("O", "Test Client Organization")?;
+    name_builder.append_entry_by_text("CN", "Test Client")?;
+    let name = name_builder.build();
+
+    let mut cert_builder = X509Builder::new()?;
+    cert_builder.set_version(2)?;
+
+    let serial = {
+        let mut bn = BigNum::new()?;
+        bn.rand(128, MsbOption::MAYBE_ZERO, false)?;
+        bn.to_asn1_integer()?
+    };
+    cert_builder.set_serial_number(&serial)?;
+
+    cert_builder.set_subject_name(&name)?;
+    cert_builder.set_issuer_name(ca_cert.subject_name())?;
+    cert_builder.set_pubkey(key)?;
+
+    let not_before = Asn1Time::days_from_now(0)?;
+    let not_after = Asn1Time::days_from_now(365)?;
+    cert_builder.set_not_before(&not_before)?;
+    cert_builder.set_not_after(&not_after)?;
+
+    // Extensions for client certificate
+    let basic_constraints = BasicConstraints::new().build()?;
+    cert_builder.append_extension(basic_constraints)?;
+
+    let key_usage = KeyUsage::new().critical().digital_signature().build()?;
+    cert_builder.append_extension(key_usage)?;
+
+    let ext_key_usage = ExtendedKeyUsage::new().client_auth().build()?;
+    cert_builder.append_extension(ext_key_usage)?;
+
+    cert_builder.sign(ca_key, MessageDigest::sha256())?;
+
+    Ok(cert_builder.build())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_certificate_generation() {
+        let certs = TestCertificates::generate().expect("Failed to generate certificates");
+
+        // Verify all files exist
+        assert!(certs.ca_cert_path.exists(), "CA cert should exist");
+        assert!(certs.ca_key_path.exists(), "CA key should exist");
+        assert!(certs.server_cert_path.exists(), "Server cert should exist");
+        assert!(certs.server_key_path.exists(), "Server key should exist");
+        assert!(certs.client_cert_path.exists(), "Client cert should exist");
+        assert!(certs.client_key_path.exists(), "Client key should exist");
+
+        // Verify files are not empty
+        assert!(
+            std::fs::metadata(&certs.ca_cert_path).unwrap().len() > 0,
+            "CA cert should not be empty"
+        );
+        assert!(
+            std::fs::metadata(&certs.server_cert_path).unwrap().len() > 0,
+            "Server cert should not be empty"
+        );
+        assert!(
+            std::fs::metadata(&certs.client_cert_path).unwrap().len() > 0,
+            "Client cert should not be empty"
+        );
+    }
+}

third_party/sglang/sgl-model-gateway/tests/common/test_config.rs

@@ -0,0 +1,336 @@
+//! Test configuration builders to reduce duplication across tests
+//!
+//! Provides pre-configured RouterConfig and MockWorkerConfig builders
+//! for common test scenarios.
+
+use smg::config::{
+    CircuitBreakerConfig, ManualAssignmentMode, PolicyConfig, RetryConfig, RouterConfig,
+};
+
+use super::mock_worker::{HealthStatus, MockWorkerConfig, WorkerType};
+
+/// Default test configuration values
+pub mod defaults {
+    pub const HOST: &str = "127.0.0.1";
+    pub const MAX_PAYLOAD_SIZE: usize = 256 * 1024 * 1024; // 256MB
+    pub const REQUEST_TIMEOUT_SECS: u64 = 600;
+    pub const WORKER_STARTUP_TIMEOUT_SECS: u64 = 5;
+    pub const WORKER_STARTUP_CHECK_INTERVAL_SECS: u64 = 1;
+    pub const MAX_CONCURRENT_REQUESTS: i32 = 64;
+    pub const QUEUE_TIMEOUT_SECS: u64 = 60;
+}
+
+/// Builder for common test RouterConfig patterns
+pub struct TestRouterConfig;
+
+impl TestRouterConfig {
+    /// Create a basic round-robin config for routing tests
+    pub fn round_robin(port: u16) -> RouterConfig {
+        RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host(defaults::HOST)
+            .port(port)
+            .max_payload_size(defaults::MAX_PAYLOAD_SIZE)
+            .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
+            .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
+            .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
+            .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
+            .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
+            .build_unchecked()
+    }
+
+    /// Create a random load balancing config
+    pub fn random(port: u16) -> RouterConfig {
+        RouterConfig::builder()
+            .regular_mode(vec![])
+            .random_policy()
+            .host(defaults::HOST)
+            .port(port)
+            .max_payload_size(defaults::MAX_PAYLOAD_SIZE)
+            .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
+            .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
+            .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
+            .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
+            .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
+            .build_unchecked()
+    }
+
+    /// Create a cache-aware config for routing tests
+    pub fn cache_aware(port: u16) -> RouterConfig {
+        RouterConfig::builder()
+            .regular_mode(vec![])
+            .cache_aware_policy(
+                0.5,  // cache_threshold
+                32,   // balance_abs_threshold
+                1.5,  // balance_rel_threshold
+                60,   // eviction_interval_secs
+                1000, // max_tree_size
+            )
+            .host(defaults::HOST)
+            .port(port)
+            .max_payload_size(defaults::MAX_PAYLOAD_SIZE)
+            .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
+            .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
+            .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
+            .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
+            .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
+            .build_unchecked()
+    }
+
+    /// Create a power-of-two config
+    pub fn power_of_two(port: u16) -> RouterConfig {
+        RouterConfig::builder()
+            .regular_mode(vec![])
+            .power_of_two_policy(5) // load_check_interval_secs
+            .host(defaults::HOST)
+            .port(port)
+            .max_payload_size(defaults::MAX_PAYLOAD_SIZE)
+            .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
+            .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
+            .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
+            .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
+            .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
+            .build_unchecked()
+    }
+
+    /// Create a manual routing config (for sticky routing tests)
+    pub fn manual(port: u16) -> RouterConfig {
+        Self::manual_with_mode(port, ManualAssignmentMode::Random)
+    }
+
+    /// Create a manual routing config with min_group assignment mode
+    pub fn manual_min_group(port: u16) -> RouterConfig {
+        Self::manual_with_mode(port, ManualAssignmentMode::MinGroup)
+    }
+
+    /// Create a manual routing config with specified assignment mode
+    pub fn manual_with_mode(port: u16, assignment_mode: ManualAssignmentMode) -> RouterConfig {
+        RouterConfig::builder()
+            .regular_mode(vec![])
+            .policy(PolicyConfig::Manual {
+                eviction_interval_secs: 60,
+                max_idle_secs: 3600,
+                assignment_mode,
+            })
+            .host(defaults::HOST)
+            .port(port)
+            .max_payload_size(defaults::MAX_PAYLOAD_SIZE)
+            .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
+            .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
+            .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
+            .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
+            .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
+            .build_unchecked()
+    }
+
+    /// Create a config with custom concurrent request limit (for rate limiting tests)
+    pub fn with_concurrency(port: u16, max_concurrent: i32) -> RouterConfig {
+        RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host(defaults::HOST)
+            .port(port)
+            .max_payload_size(defaults::MAX_PAYLOAD_SIZE)
+            .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
+            .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
+            .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
+            .max_concurrent_requests(max_concurrent)
+            .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
+            .build_unchecked()
+    }
+
+    /// Create a config with custom payload size limit
+    pub fn with_payload_limit(port: u16, max_payload_size: usize) -> RouterConfig {
+        RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host(defaults::HOST)
+            .port(port)
+            .max_payload_size(max_payload_size)
+            .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
+            .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
+            .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
+            .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
+            .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
+            .build_unchecked()
+    }
+
+    /// Create a config with short timeouts (for timeout/retry tests)
+    pub fn with_short_timeouts(port: u16) -> RouterConfig {
+        RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host(defaults::HOST)
+            .port(port)
+            .max_payload_size(defaults::MAX_PAYLOAD_SIZE)
+            .request_timeout_secs(5)
+            .worker_startup_timeout_secs(2)
+            .worker_startup_check_interval_secs(1)
+            .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
+            .queue_timeout_secs(5)
+            .build_unchecked()
+    }
+
+    /// Create a round-robin config with retry settings
+    pub fn round_robin_with_retry(port: u16, retry_config: RetryConfig) -> RouterConfig {
+        RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host(defaults::HOST)
+            .port(port)
+            .max_payload_size(defaults::MAX_PAYLOAD_SIZE)
+            .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
+            .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
+            .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
+            .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
+            .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
+            .retry_config(retry_config)
+            .build_unchecked()
+    }
+
+    /// Create a round-robin config with circuit breaker
+    pub fn round_robin_with_circuit_breaker(
+        port: u16,
+        circuit_breaker: CircuitBreakerConfig,
+    ) -> RouterConfig {
+        RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host(defaults::HOST)
+            .port(port)
+            .max_payload_size(defaults::MAX_PAYLOAD_SIZE)
+            .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
+            .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
+            .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
+            .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
+            .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
+            .circuit_breaker_config(circuit_breaker)
+            .build_unchecked()
+    }
+
+    /// Create a round-robin config with both retry and circuit breaker
+    pub fn round_robin_with_reliability(
+        port: u16,
+        retry_config: RetryConfig,
+        circuit_breaker: CircuitBreakerConfig,
+    ) -> RouterConfig {
+        RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host(defaults::HOST)
+            .port(port)
+            .max_payload_size(defaults::MAX_PAYLOAD_SIZE)
+            .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
+            .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
+            .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
+            .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
+            .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
+            .retry_config(retry_config)
+            .circuit_breaker_config(circuit_breaker)
+            .build_unchecked()
+    }
+}
+
+/// Builder for common MockWorkerConfig patterns
+pub struct TestWorkerConfig;
+
+impl TestWorkerConfig {
+    /// Create a healthy worker config
+    pub fn healthy(port: u16) -> MockWorkerConfig {
+        MockWorkerConfig {
+            port,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }
+    }
+
+    /// Create multiple healthy workers with sequential ports
+    pub fn healthy_workers(start_port: u16, count: u16) -> Vec<MockWorkerConfig> {
+        (0..count).map(|i| Self::healthy(start_port + i)).collect()
+    }
+
+    /// Create an unhealthy worker config
+    pub fn unhealthy(port: u16) -> MockWorkerConfig {
+        MockWorkerConfig {
+            port,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Unhealthy,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }
+    }
+
+    /// Create a slow worker config (for timeout tests)
+    pub fn slow(port: u16, delay_ms: u64) -> MockWorkerConfig {
+        MockWorkerConfig {
+            port,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: delay_ms,
+            fail_rate: 0.0,
+        }
+    }
+
+    /// Create multiple slow workers with sequential ports
+    pub fn slow_workers(start_port: u16, count: u16, delay_ms: u64) -> Vec<MockWorkerConfig> {
+        (0..count)
+            .map(|i| Self::slow(start_port + i, delay_ms))
+            .collect()
+    }
+
+    /// Create a flaky worker config (for retry/fault tolerance tests)
+    pub fn flaky(port: u16, fail_rate: f32) -> MockWorkerConfig {
+        MockWorkerConfig {
+            port,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate,
+        }
+    }
+
+    /// Create a decode worker config (for PD routing tests)
+    pub fn decode(port: u16) -> MockWorkerConfig {
+        MockWorkerConfig {
+            port,
+            worker_type: WorkerType::Decode,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }
+    }
+
+    /// Create a prefill worker config (for PD routing tests)
+    pub fn prefill(port: u16) -> MockWorkerConfig {
+        MockWorkerConfig {
+            port,
+            worker_type: WorkerType::Prefill,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_round_robin_config() {
+        let config = TestRouterConfig::round_robin(3000);
+        assert_eq!(config.port, 3000);
+    }
+
+    #[test]
+    fn test_healthy_workers() {
+        let workers = TestWorkerConfig::healthy_workers(8000, 3);
+        assert_eq!(workers.len(), 3);
+        assert_eq!(workers[0].port, 8000);
+        assert_eq!(workers[1].port, 8001);
+        assert_eq!(workers[2].port, 8002);
+    }
+}

third_party/sglang/sgl-model-gateway/tests/common/tls_mock_worker.rs

@@ -0,0 +1,370 @@
+// TLS-enabled mock worker for mTLS integration tests
+#![allow(dead_code)]
+
+use std::{
+    net::SocketAddr,
+    path::Path,
+    sync::{Arc, Once},
+    time::{SystemTime, UNIX_EPOCH},
+};
+
+use axum::{
+    extract::{Json, State},
+    http::StatusCode,
+    response::{IntoResponse, Response},
+    routing::{get, post},
+    Router,
+};
+use axum_server::tls_rustls::RustlsConfig;
+use rustls::{server::WebPkiClientVerifier, RootCertStore};
+use serde_json::json;
+use tokio::sync::RwLock;
+
+// Ensure crypto provider is installed exactly once
+static CRYPTO_PROVIDER_INIT: Once = Once::new();
+
+fn ensure_crypto_provider() {
+    CRYPTO_PROVIDER_INIT.call_once(|| {
+        let _ = rustls::crypto::ring::default_provider().install_default();
+    });
+}
+
+/// Configuration for TLS mock worker behavior
+#[derive(Clone)]
+pub struct TlsMockWorkerConfig {
+    pub port: u16,
+    /// Require client certificate (mTLS) or just server TLS
+    pub require_client_cert: bool,
+    /// Response delay in milliseconds
+    pub response_delay_ms: u64,
+    /// Fail rate (0.0 - 1.0)
+    pub fail_rate: f32,
+}
+
+impl Default for TlsMockWorkerConfig {
+    fn default() -> Self {
+        Self {
+            port: 0,
+            require_client_cert: true,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }
+    }
+}
+
+/// TLS-enabled mock worker server for mTLS testing
+pub struct TlsMockWorker {
+    config: Arc<RwLock<TlsMockWorkerConfig>>,
+    shutdown_handle: Option<tokio::task::JoinHandle<()>>,
+    shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
+}
+
+impl TlsMockWorker {
+    pub fn new(config: TlsMockWorkerConfig) -> Self {
+        Self {
+            config: Arc::new(RwLock::new(config)),
+            shutdown_handle: None,
+            shutdown_tx: None,
+        }
+    }
+
+    /// Start the TLS mock worker server
+    ///
+    /// # Arguments
+    /// * `server_cert_path` - Path to server certificate PEM file
+    /// * `server_key_path` - Path to server private key PEM file
+    /// * `ca_cert_path` - Path to CA certificate for client verification (optional for TLS-only mode)
+    pub async fn start(
+        &mut self,
+        server_cert_path: &Path,
+        server_key_path: &Path,
+        ca_cert_path: Option<&Path>,
+    ) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
+        // Ensure crypto provider is installed before using rustls
+        ensure_crypto_provider();
+
+        let config = self.config.clone();
+        let port = config.read().await.port;
+        let require_client_cert = config.read().await.require_client_cert;
+
+        // If port is 0, find an available port
+        let port = if port == 0 {
+            let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
+            let port = listener.local_addr()?.port();
+            drop(listener);
+            config.write().await.port = port;
+            port
+        } else {
+            port
+        };
+
+        let app = Router::new()
+            .route("/health", get(health_handler))
+            .route("/health_generate", get(health_generate_handler))
+            .route("/server_info", get(server_info_handler))
+            .route("/generate", post(generate_handler))
+            .route("/v1/chat/completions", post(chat_completions_handler))
+            .with_state(config);
+
+        let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel::<()>();
+        self.shutdown_tx = Some(shutdown_tx);
+
+        // Build TLS configuration
+        let rustls_config = if require_client_cert {
+            // mTLS: require client certificate
+            let ca_cert_path = ca_cert_path.ok_or("CA cert path required for mTLS")?;
+            build_mtls_config(server_cert_path, server_key_path, ca_cert_path).await?
+        } else {
+            // TLS only: no client cert required
+            build_tls_config(server_cert_path, server_key_path).await?
+        };
+
+        let addr = SocketAddr::from(([127, 0, 0, 1], port));
+
+        // Spawn the server in a separate task
+        let handle = tokio::spawn(async move {
+            let server =
+                axum_server::bind_rustls(addr, rustls_config).serve(app.into_make_service());
+
+            tokio::select! {
+                result = server => {
+                    if let Err(e) = result {
+                        eprintln!("TLS Server error: {}", e);
+                    }
+                }
+                _ = &mut shutdown_rx => {
+                    // Graceful shutdown
+                }
+            }
+        });
+
+        self.shutdown_handle = Some(handle);
+
+        // Wait for the server to start
+        tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
+
+        let url = format!("https://127.0.0.1:{}", port);
+        Ok(url)
+    }
+
+    /// Stop the TLS mock worker server
+    pub async fn stop(&mut self) {
+        if let Some(shutdown_tx) = self.shutdown_tx.take() {
+            let _ = shutdown_tx.send(());
+        }
+
+        if let Some(handle) = self.shutdown_handle.take() {
+            let _ = tokio::time::timeout(tokio::time::Duration::from_secs(5), handle).await;
+        }
+    }
+}
+
+impl Drop for TlsMockWorker {
+    fn drop(&mut self) {
+        if let Some(shutdown_tx) = self.shutdown_tx.take() {
+            let _ = shutdown_tx.send(());
+        }
+    }
+}
+
+/// Build TLS config for server-only TLS (no client cert required)
+async fn build_tls_config(
+    cert_path: &Path,
+    key_path: &Path,
+) -> Result<RustlsConfig, Box<dyn std::error::Error + Send + Sync>> {
+    let config = RustlsConfig::from_pem_file(cert_path, key_path).await?;
+    Ok(config)
+}
+
+/// Build mTLS config requiring client certificate
+async fn build_mtls_config(
+    cert_path: &Path,
+    key_path: &Path,
+    ca_cert_path: &Path,
+) -> Result<RustlsConfig, Box<dyn std::error::Error + Send + Sync>> {
+    use std::io::BufReader;
+
+    use rustls_pemfile::{certs, pkcs8_private_keys};
+
+    // Read server certificate
+    let cert_file = std::fs::File::open(cert_path)?;
+    let mut reader = BufReader::new(cert_file);
+    let cert_chain: Vec<_> = certs(&mut reader).filter_map(|r| r.ok()).collect();
+
+    // Read server private key
+    let key_file = std::fs::File::open(key_path)?;
+    let mut reader = BufReader::new(key_file);
+    let private_key = pkcs8_private_keys(&mut reader)
+        .next()
+        .ok_or("No private key found")??;
+
+    // Read CA certificate for client verification
+    let ca_file = std::fs::File::open(ca_cert_path)?;
+    let mut reader = BufReader::new(ca_file);
+    let ca_certs: Vec<_> = certs(&mut reader).filter_map(|r| r.ok()).collect();
+
+    // Build root certificate store for client verification
+    let mut root_store = RootCertStore::empty();
+    for cert in ca_certs {
+        root_store.add(cert)?;
+    }
+
+    // Create client certificate verifier
+    let client_verifier = WebPkiClientVerifier::builder(Arc::new(root_store))
+        .build()
+        .map_err(|e| format!("Failed to build client verifier: {}", e))?;
+
+    // Build server config with client verification
+    let server_config = rustls::ServerConfig::builder()
+        .with_client_cert_verifier(client_verifier)
+        .with_single_cert(cert_chain, private_key.into())
+        .map_err(|e| format!("Failed to build server config: {}", e))?;
+
+    Ok(RustlsConfig::from_config(Arc::new(server_config)))
+}
+
+// Handler implementations (simplified versions of mock_worker handlers)
+
+async fn should_fail(config: &TlsMockWorkerConfig) -> bool {
+    rand::random::<f32>() < config.fail_rate
+}
+
+async fn health_handler(State(config): State<Arc<RwLock<TlsMockWorkerConfig>>>) -> Response {
+    let config = config.read().await;
+
+    if should_fail(&config).await {
+        return (
+            StatusCode::INTERNAL_SERVER_ERROR,
+            Json(json!({ "error": "Random failure" })),
+        )
+            .into_response();
+    }
+
+    Json(json!({
+        "status": "healthy",
+        "timestamp": SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(),
+        "tls_enabled": true
+    }))
+    .into_response()
+}
+
+async fn health_generate_handler(
+    State(config): State<Arc<RwLock<TlsMockWorkerConfig>>>,
+) -> Response {
+    let config = config.read().await;
+
+    if should_fail(&config).await {
+        return (
+            StatusCode::INTERNAL_SERVER_ERROR,
+            Json(json!({ "error": "Random failure" })),
+        )
+            .into_response();
+    }
+
+    Json(json!({
+        "status": "ok",
+        "queue_length": 0,
+        "processing_time_ms": config.response_delay_ms
+    }))
+    .into_response()
+}
+
+async fn server_info_handler(State(config): State<Arc<RwLock<TlsMockWorkerConfig>>>) -> Response {
+    let config = config.read().await;
+
+    if should_fail(&config).await {
+        return (
+            StatusCode::INTERNAL_SERVER_ERROR,
+            Json(json!({ "error": "Random failure" })),
+        )
+            .into_response();
+    }
+
+    Json(json!({
+        "model_path": "mock-tls-model",
+        "port": config.port,
+        "host": "127.0.0.1",
+        "tls_enabled": true,
+        "version": "0.3.0"
+    }))
+    .into_response()
+}
+
+async fn generate_handler(
+    State(config): State<Arc<RwLock<TlsMockWorkerConfig>>>,
+    Json(_payload): Json<serde_json::Value>,
+) -> Response {
+    let config = config.read().await;
+
+    if should_fail(&config).await {
+        return (
+            StatusCode::INTERNAL_SERVER_ERROR,
+            Json(json!({ "error": "Random failure" })),
+        )
+            .into_response();
+    }
+
+    if config.response_delay_ms > 0 {
+        tokio::time::sleep(tokio::time::Duration::from_millis(config.response_delay_ms)).await;
+    }
+
+    Json(json!({
+        "text": "This is a mock TLS response.",
+        "meta_info": {
+            "prompt_tokens": 10,
+            "completion_tokens": 5,
+            "tls_verified": true
+        }
+    }))
+    .into_response()
+}
+
+async fn chat_completions_handler(
+    State(config): State<Arc<RwLock<TlsMockWorkerConfig>>>,
+    Json(_payload): Json<serde_json::Value>,
+) -> Response {
+    let config = config.read().await;
+
+    if should_fail(&config).await {
+        return (
+            StatusCode::INTERNAL_SERVER_ERROR,
+            Json(json!({
+                "error": {
+                    "message": "Random failure",
+                    "type": "internal_error"
+                }
+            })),
+        )
+            .into_response();
+    }
+
+    if config.response_delay_ms > 0 {
+        tokio::time::sleep(tokio::time::Duration::from_millis(config.response_delay_ms)).await;
+    }
+
+    let timestamp = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_secs();
+
+    Json(json!({
+        "id": format!("chatcmpl-{}", uuid::Uuid::new_v4()),
+        "object": "chat.completion",
+        "created": timestamp,
+        "model": "mock-tls-model",
+        "choices": [{
+            "index": 0,
+            "message": {
+                "role": "assistant",
+                "content": "This is a mock TLS chat response."
+            },
+            "finish_reason": "stop"
+        }],
+        "usage": {
+            "prompt_tokens": 10,
+            "completion_tokens": 5,
+            "total_tokens": 15
+        }
+    }))
+    .into_response()
+}

third_party/sglang/sgl-model-gateway/tests/inflight_tracker_test.rs

@@ -0,0 +1,148 @@
+mod common;
+
+use std::time::Duration;
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use common::{
+    mock_worker::{HealthStatus, MockWorkerConfig, WorkerType},
+    AppTestContext,
+};
+use serde_json::json;
+use tower::ServiceExt;
+
+#[tokio::test]
+async fn test_multiple_concurrent_requests_tracking() {
+    let ctx = AppTestContext::new(vec![MockWorkerConfig {
+        port: 19002,
+        worker_type: WorkerType::Regular,
+        health_status: HealthStatus::Healthy,
+        response_delay_ms: 50,
+        fail_rate: 0.0,
+    }])
+    .await;
+
+    let tracker = ctx.app_context.inflight_tracker.clone();
+
+    let mut handles = vec![];
+    for i in 0..5 {
+        let app = ctx.create_app().await;
+        handles.push(tokio::spawn(async move {
+            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();
+
+            app.oneshot(req).await.unwrap()
+        }));
+    }
+
+    for handle in handles {
+        let resp = handle.await.unwrap();
+        assert_eq!(resp.status(), StatusCode::OK);
+    }
+
+    assert!(tracker.is_empty());
+
+    ctx.shutdown().await;
+}
+
+#[tokio::test]
+async fn test_inflight_request_appears_in_bucket() {
+    let ctx = AppTestContext::new(vec![MockWorkerConfig {
+        port: 19004,
+        worker_type: WorkerType::Regular,
+        health_status: HealthStatus::Healthy,
+        response_delay_ms: 2000,
+        fail_rate: 0.0,
+    }])
+    .await;
+
+    let tracker = &ctx.app_context.inflight_tracker;
+    assert!(tracker.is_empty(), "Tracker should start empty");
+
+    let app = ctx.create_app().await;
+    let payload = json!({
+        "text": "Long running request",
+        "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 tracker_clone = ctx.app_context.inflight_tracker.clone();
+    let response_future = tokio::spawn(async move { app.oneshot(req).await });
+
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    let inflight_count = tracker_clone.len();
+    assert!(
+        inflight_count > 0,
+        "Should have at least one in-flight request, got {}",
+        inflight_count
+    );
+
+    let buckets = tracker_clone.compute_bucket_counts();
+    assert!(buckets[0] > 0, "first bucket (<=30s) should have requests");
+
+    let resp = response_future.await.unwrap().unwrap();
+    assert_eq!(resp.status(), StatusCode::OK);
+
+    tokio::time::sleep(Duration::from_millis(50)).await;
+    assert!(
+        tracker_clone.is_empty(),
+        "Request should be deregistered after completion"
+    );
+
+    ctx.shutdown().await;
+}
+
+#[tokio::test]
+async fn test_failed_request_still_deregisters() {
+    let ctx = AppTestContext::new(vec![MockWorkerConfig {
+        port: 19003,
+        worker_type: WorkerType::Regular,
+        health_status: HealthStatus::Healthy,
+        response_delay_ms: 0,
+        fail_rate: 1.0,
+    }])
+    .await;
+
+    let tracker = &ctx.app_context.inflight_tracker;
+    assert!(tracker.is_empty());
+
+    let app = ctx.create_app().await;
+
+    let payload = json!({
+        "text": "This should fail",
+        "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::INTERNAL_SERVER_ERROR);
+
+    assert!(tracker.is_empty());
+
+    ctx.shutdown().await;
+}

third_party/sglang/sgl-model-gateway/tests/load_guard_raii_test.rs

@@ -0,0 +1,212 @@
+//! Tests for WorkerLoadGuard RAII pattern with response body attachment
+//!
+//! These tests verify that load guards properly decrement worker load when:
+//! - Response body is fully consumed
+//! - Response body is dropped (client disconnect simulation)
+//! - Multiple guards are attached (dual prefill/decode workers)
+
+use std::sync::Arc;
+
+use axum::{body::Body, response::Response};
+use bytes::Bytes;
+use futures_util::StreamExt;
+use http_body_util::BodyExt;
+use smg::core::{AttachedBody, BasicWorkerBuilder, Worker, WorkerLoadGuard};
+use tokio::sync::mpsc;
+use tokio_stream::wrappers::UnboundedReceiverStream;
+
+/// Helper to create an SSE streaming response
+fn create_sse_response(rx: mpsc::UnboundedReceiver<Bytes>) -> Response {
+    let stream = UnboundedReceiverStream::new(rx).map(Ok::<_, std::io::Error>);
+    let body = Body::from_stream(stream);
+    Response::new(body)
+}
+
+/// Helper to create a test worker
+fn create_test_worker() -> Arc<dyn Worker> {
+    Arc::new(BasicWorkerBuilder::new("http://localhost:8000").build())
+}
+
+#[tokio::test]
+async fn test_guard_dropped_when_response_body_consumed() {
+    let worker = create_test_worker();
+    assert_eq!(worker.load(), 0);
+
+    // Create a simple response with some data
+    let body = Body::from("Hello, World!");
+    let response = Response::new(body);
+
+    // Attach guard
+    let guard = WorkerLoadGuard::new(worker.clone(), None);
+    assert_eq!(worker.load(), 1);
+
+    let guarded_response = AttachedBody::wrap_response(response, guard);
+
+    // Load should still be 1 (guard is in the body)
+    assert_eq!(worker.load(), 1);
+
+    // Consume the response body
+    let body = guarded_response.into_body();
+    let _bytes = body.collect().await.unwrap().to_bytes();
+
+    // After consuming, guard should be dropped, load should be 0
+    assert_eq!(worker.load(), 0);
+}
+
+#[tokio::test]
+async fn test_guard_dropped_when_response_dropped_without_consumption() {
+    let worker = create_test_worker();
+    assert_eq!(worker.load(), 0);
+
+    {
+        let body = Body::from("Hello, World!");
+        let response = Response::new(body);
+
+        let guard = WorkerLoadGuard::new(worker.clone(), None);
+        assert_eq!(worker.load(), 1);
+
+        let _guarded_response = AttachedBody::wrap_response(response, guard);
+
+        // Load is still 1
+        assert_eq!(worker.load(), 1);
+
+        // Response goes out of scope here
+    }
+
+    // After response is dropped, guard should be dropped, load should be 0
+    assert_eq!(worker.load(), 0);
+}
+
+#[tokio::test]
+async fn test_streaming_guard_dropped_when_stream_ends() {
+    let worker = create_test_worker();
+    assert_eq!(worker.load(), 0);
+
+    // Create a channel for SSE streaming
+    let (tx, rx) = mpsc::unbounded_channel::<Bytes>();
+
+    let response = create_sse_response(rx);
+    let guard = WorkerLoadGuard::new(worker.clone(), None);
+    assert_eq!(worker.load(), 1);
+
+    let guarded_response = AttachedBody::wrap_response(response, guard);
+
+    // Spawn a task to consume the response
+    let worker_clone = worker.clone();
+    let consume_task = tokio::spawn(async move {
+        {
+            let mut body = guarded_response.into_body();
+            while let Some(result) = body.frame().await {
+                if result.is_err() {
+                    break;
+                }
+            }
+            // Body is still in scope here, guard not dropped yet
+        }
+        // Body dropped here, guard should be dropped
+        assert_eq!(worker_clone.load(), 0);
+    });
+
+    // Send some data
+    tx.send(Bytes::from("data: chunk1\n\n")).unwrap();
+    tx.send(Bytes::from("data: chunk2\n\n")).unwrap();
+
+    // Load should still be 1 while streaming
+    tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
+    assert_eq!(worker.load(), 1);
+
+    // Close the sender to end the stream
+    drop(tx);
+
+    // Wait for consumer to finish
+    consume_task.await.unwrap();
+
+    // Load should now be 0
+    assert_eq!(worker.load(), 0);
+}
+
+#[tokio::test]
+async fn test_streaming_guard_dropped_on_client_disconnect() {
+    let worker = create_test_worker();
+    assert_eq!(worker.load(), 0);
+
+    let (tx, rx) = mpsc::unbounded_channel::<Bytes>();
+
+    let response = create_sse_response(rx);
+    let guard = WorkerLoadGuard::new(worker.clone(), None);
+    assert_eq!(worker.load(), 1);
+
+    let guarded_response = AttachedBody::wrap_response(response, guard);
+
+    // Start consuming but drop early (simulate client disconnect)
+    {
+        let mut body = guarded_response.into_body();
+
+        // Read one frame
+        tx.send(Bytes::from("data: chunk1\n\n")).unwrap();
+        let _ = body.frame().await;
+
+        // Load still 1
+        assert_eq!(worker.load(), 1);
+
+        // Body dropped here (simulating client disconnect)
+    }
+
+    // Guard should be dropped when body is dropped
+    assert_eq!(worker.load(), 0);
+
+    // tx is still open but no one is listening
+    drop(tx);
+}
+
+#[tokio::test]
+async fn test_multiple_guards_all_dropped() {
+    let worker1 = create_test_worker();
+    let worker2 = create_test_worker();
+    assert_eq!(worker1.load(), 0);
+    assert_eq!(worker2.load(), 0);
+
+    {
+        let body = Body::from("Hello");
+        let response = Response::new(body);
+
+        // Create guards for both workers (simulates dual prefill/decode)
+        let guard1 = WorkerLoadGuard::new(worker1.clone(), None);
+        let guard2 = WorkerLoadGuard::new(worker2.clone(), None);
+        assert_eq!(worker1.load(), 1);
+        assert_eq!(worker2.load(), 1);
+
+        let _response = AttachedBody::wrap_response(response, vec![guard1, guard2]);
+
+        // Both loads are 1
+        assert_eq!(worker1.load(), 1);
+        assert_eq!(worker2.load(), 1);
+    }
+
+    // Both guards dropped when response goes out of scope
+    assert_eq!(worker1.load(), 0);
+    assert_eq!(worker2.load(), 0);
+}
+
+#[tokio::test]
+async fn test_guard_with_empty_body() {
+    let worker = create_test_worker();
+    assert_eq!(worker.load(), 0);
+
+    {
+        let body = Body::empty();
+        let response = Response::new(body);
+
+        let guard = WorkerLoadGuard::new(worker.clone(), None);
+        assert_eq!(worker.load(), 1);
+
+        let guarded_response = AttachedBody::wrap_response(response, guard);
+
+        // Consume empty body
+        let body = guarded_response.into_body();
+        let bytes = body.collect().await.unwrap().to_bytes();
+        assert!(bytes.is_empty());
+    }
+
+    assert_eq!(worker.load(), 0);
+}

third_party/sglang/sgl-model-gateway/tests/mcp_test.rs

@@ -0,0 +1,543 @@
+// This test suite validates the complete MCP implementation against the
+// functionality required for SGLang responses API integration.
+//
+// - Core MCP server functionality
+// - Tool session management (individual and multi-tool)
+// - Tool execution and error handling
+// - Schema adaptation and validation
+// - Mock server integration for reliable testing
+
+mod common;
+
+use std::collections::HashMap;
+
+use common::mock_mcp_server::MockMCPServer;
+use serde_json::json;
+use smg_mcp::{error::McpError, McpConfig, McpManager, McpServerConfig, McpTransport};
+
+/// Create a new mock server for testing (each test gets its own)
+async fn create_mock_server() -> MockMCPServer {
+    MockMCPServer::start()
+        .await
+        .expect("Failed to start mock MCP server")
+}
+
+// Core MCP Server Tests
+
+#[tokio::test]
+async fn test_mcp_server_initialization() {
+    let config = McpConfig {
+        servers: vec![],
+        pool: Default::default(),
+        proxy: None,
+        warmup: Vec::new(),
+        inventory: Default::default(),
+    };
+
+    // Should succeed but with no connected servers (empty config is allowed)
+    let result = McpManager::with_defaults(config).await;
+    assert!(result.is_ok(), "Should succeed with empty config");
+
+    let manager = result.unwrap();
+    let servers = manager.list_servers();
+    assert_eq!(servers.len(), 0, "Should have no servers");
+    let tools = manager.list_tools();
+    assert_eq!(tools.len(), 0, "Should have no tools");
+}
+
+#[tokio::test]
+async fn test_server_connection_with_mock() {
+    let mock_server = create_mock_server().await;
+
+    let config = McpConfig {
+        servers: vec![McpServerConfig {
+            name: "mock_server".to_string(),
+            transport: McpTransport::Streamable {
+                url: mock_server.url(),
+                token: None,
+            },
+            proxy: None,
+            required: false,
+        }],
+        pool: Default::default(),
+        proxy: None,
+        warmup: Vec::new(),
+        inventory: Default::default(),
+    };
+
+    let result = McpManager::with_defaults(config).await;
+    assert!(result.is_ok(), "Should connect to mock server");
+
+    let manager = result.unwrap();
+
+    let servers = manager.list_servers();
+    assert_eq!(servers.len(), 1);
+    assert!(servers.contains(&"mock_server".to_string()));
+
+    let tools = manager.list_tools();
+    assert_eq!(tools.len(), 2, "Should have 2 tools from mock server");
+
+    assert!(manager.has_tool("brave_web_search"));
+    assert!(manager.has_tool("brave_local_search"));
+
+    manager.shutdown().await;
+}
+
+#[tokio::test]
+async fn test_tool_availability_checking() {
+    let mock_server = create_mock_server().await;
+
+    let config = McpConfig {
+        servers: vec![McpServerConfig {
+            name: "mock_server".to_string(),
+            transport: McpTransport::Streamable {
+                url: mock_server.url(),
+                token: None,
+            },
+            proxy: None,
+            required: false,
+        }],
+        pool: Default::default(),
+        proxy: None,
+        warmup: Vec::new(),
+        inventory: Default::default(),
+    };
+
+    let manager = McpManager::with_defaults(config).await.unwrap();
+
+    let test_tools = vec!["brave_web_search", "brave_local_search", "calculator"];
+    for tool in test_tools {
+        let available = manager.has_tool(tool);
+        match tool {
+            "brave_web_search" | "brave_local_search" => {
+                assert!(
+                    available,
+                    "Tool {} should be available from mock server",
+                    tool
+                );
+            }
+            "calculator" => {
+                assert!(
+                    !available,
+                    "Tool {} should not be available from mock server",
+                    tool
+                );
+            }
+            _ => {}
+        }
+    }
+
+    manager.shutdown().await;
+}
+
+#[tokio::test]
+async fn test_multi_server_connection() {
+    let mock_server1 = create_mock_server().await;
+    let mock_server2 = create_mock_server().await;
+
+    let config = McpConfig {
+        servers: vec![
+            McpServerConfig {
+                name: "mock_server_1".to_string(),
+                transport: McpTransport::Streamable {
+                    url: mock_server1.url(),
+                    token: None,
+                },
+                proxy: None,
+                required: false,
+            },
+            McpServerConfig {
+                name: "mock_server_2".to_string(),
+                transport: McpTransport::Streamable {
+                    url: mock_server2.url(),
+                    token: None,
+                },
+                proxy: None,
+                required: false,
+            },
+        ],
+        pool: Default::default(),
+        proxy: None,
+        warmup: Vec::new(),
+        inventory: Default::default(),
+    };
+
+    // Note: This will fail to connect to both servers in the current implementation
+    // since they return the same tools. The manager will connect to the first one.
+    let result = McpManager::with_defaults(config).await;
+
+    if let Ok(manager) = result {
+        let servers = manager.list_servers();
+        assert!(!servers.is_empty(), "Should have at least one server");
+
+        let tools = manager.list_tools();
+        assert!(tools.len() >= 2, "Should have tools from servers");
+
+        manager.shutdown().await;
+    }
+}
+
+#[tokio::test]
+async fn test_tool_execution_with_mock() {
+    let mock_server = create_mock_server().await;
+
+    let config = McpConfig {
+        servers: vec![McpServerConfig {
+            name: "mock_server".to_string(),
+            transport: McpTransport::Streamable {
+                url: mock_server.url(),
+                token: None,
+            },
+            proxy: None,
+            required: false,
+        }],
+        pool: Default::default(),
+        proxy: None,
+        warmup: Vec::new(),
+        inventory: Default::default(),
+    };
+
+    let manager = McpManager::with_defaults(config).await.unwrap();
+
+    let result = manager
+        .call_tool(
+            "brave_web_search",
+            Some(
+                json!({
+                    "query": "rust programming",
+                    "count": 1
+                })
+                .as_object()
+                .unwrap()
+                .clone(),
+            ),
+        )
+        .await;
+
+    assert!(
+        result.is_ok(),
+        "Tool execution should succeed with mock server"
+    );
+
+    let response = result.unwrap();
+    assert!(!response.content.is_empty(), "Should have content");
+
+    // Check the content
+    if let rmcp::model::RawContent::Text(text) = &response.content[0].raw {
+        assert!(text
+            .text
+            .contains("Mock search results for: rust programming"));
+    } else {
+        panic!("Expected text content");
+    }
+
+    manager.shutdown().await;
+}
+
+#[tokio::test]
+async fn test_concurrent_tool_execution() {
+    let mock_server = create_mock_server().await;
+
+    let config = McpConfig {
+        servers: vec![McpServerConfig {
+            name: "mock_server".to_string(),
+            transport: McpTransport::Streamable {
+                url: mock_server.url(),
+                token: None,
+            },
+            proxy: None,
+            required: false,
+        }],
+        pool: Default::default(),
+        proxy: None,
+        warmup: Vec::new(),
+        inventory: Default::default(),
+    };
+
+    let manager = McpManager::with_defaults(config).await.unwrap();
+
+    // Execute tools sequentially (true concurrent execution would require Arc<Mutex>)
+    let tool_calls = vec![
+        ("brave_web_search", json!({"query": "test1"})),
+        ("brave_local_search", json!({"query": "test2"})),
+    ];
+
+    for (tool_name, args) in tool_calls {
+        let result = manager
+            .call_tool(tool_name, Some(args.as_object().unwrap().clone()))
+            .await;
+
+        assert!(result.is_ok(), "Tool {} should succeed", tool_name);
+        let response = result.unwrap();
+        assert!(!response.content.is_empty(), "Should have content");
+    }
+
+    manager.shutdown().await;
+}
+
+// Error Handling Tests
+
+#[tokio::test]
+async fn test_tool_execution_errors() {
+    let mock_server = create_mock_server().await;
+
+    let config = McpConfig {
+        servers: vec![McpServerConfig {
+            name: "mock_server".to_string(),
+            transport: McpTransport::Streamable {
+                url: mock_server.url(),
+                token: None,
+            },
+            proxy: None,
+            required: false,
+        }],
+        pool: Default::default(),
+        proxy: None,
+        warmup: Vec::new(),
+        inventory: Default::default(),
+    };
+
+    let manager = McpManager::with_defaults(config).await.unwrap();
+
+    // Try to call unknown tool
+    let result = manager
+        .call_tool("unknown_tool", Some(serde_json::Map::new()))
+        .await;
+    assert!(result.is_err(), "Should fail for unknown tool");
+
+    match result.unwrap_err() {
+        McpError::ToolNotFound(name) => {
+            assert_eq!(name, "unknown_tool");
+        }
+        _ => panic!("Expected ToolNotFound error"),
+    }
+
+    manager.shutdown().await;
+}
+
+#[tokio::test]
+async fn test_connection_without_server() {
+    let config = McpConfig {
+        servers: vec![McpServerConfig {
+            name: "nonexistent".to_string(),
+            transport: McpTransport::Stdio {
+                command: "/nonexistent/command".to_string(),
+                args: vec![],
+                envs: HashMap::new(),
+            },
+            proxy: None,
+            required: false,
+        }],
+        pool: Default::default(),
+        proxy: None,
+        warmup: Vec::new(),
+        inventory: Default::default(),
+    };
+
+    let result = McpManager::with_defaults(config).await;
+    // Manager succeeds but no servers are connected (errors are logged)
+    assert!(
+        result.is_ok(),
+        "Manager should succeed even if servers fail to connect"
+    );
+
+    let manager = result.unwrap();
+    let servers = manager.list_servers();
+    assert_eq!(servers.len(), 0, "Should have no connected servers");
+}
+
+// Schema Validation Tests
+
+#[tokio::test]
+async fn test_tool_info_structure() {
+    let mock_server = create_mock_server().await;
+
+    let config = McpConfig {
+        servers: vec![McpServerConfig {
+            name: "mock_server".to_string(),
+            transport: McpTransport::Streamable {
+                url: mock_server.url(),
+                token: None,
+            },
+            proxy: None,
+            required: false,
+        }],
+        pool: Default::default(),
+        proxy: None,
+        warmup: Vec::new(),
+        inventory: Default::default(),
+    };
+
+    let manager = McpManager::with_defaults(config).await.unwrap();
+
+    let tools = manager.list_tools();
+    let brave_search = tools
+        .iter()
+        .find(|t| t.name.as_ref() == "brave_web_search")
+        .expect("Should have brave_web_search tool");
+
+    assert_eq!(brave_search.name.as_ref(), "brave_web_search");
+    assert!(brave_search
+        .description
+        .as_ref()
+        .map(|d| d.contains("Mock web search"))
+        .unwrap_or(false));
+    // Note: server information is now maintained separately in the inventory,
+    // not in the Tool type itself
+    assert!(!brave_search.input_schema.is_empty());
+}
+
+// SSE Parsing Tests (simplified since we don't expose parse_sse_event)
+
+#[tokio::test]
+async fn test_sse_connection() {
+    // This tests that SSE configuration is properly handled even when connection fails
+    let config = McpConfig {
+        servers: vec![McpServerConfig {
+            name: "sse_test".to_string(),
+            transport: McpTransport::Stdio {
+                command: "/nonexistent/sse/server".to_string(),
+                args: vec!["--sse".to_string()],
+                envs: HashMap::new(),
+            },
+            proxy: None,
+            required: false,
+        }],
+        pool: Default::default(),
+        proxy: None,
+        warmup: Vec::new(),
+        inventory: Default::default(),
+    };
+
+    // Manager succeeds but no servers are connected (errors are logged)
+    let result = McpManager::with_defaults(config).await;
+    assert!(
+        result.is_ok(),
+        "Manager should succeed even if SSE server fails to connect"
+    );
+
+    let manager = result.unwrap();
+    let servers = manager.list_servers();
+    assert_eq!(servers.len(), 0, "Should have no connected servers");
+}
+
+// Connection Type Tests
+
+#[tokio::test]
+async fn test_transport_types() {
+    // HTTP/Streamable transport
+    let http_config = McpServerConfig {
+        name: "http_server".to_string(),
+        transport: McpTransport::Streamable {
+            url: "http://localhost:8080/mcp".to_string(),
+            token: Some("auth_token".to_string()),
+        },
+        proxy: None,
+        required: false,
+    };
+    assert_eq!(http_config.name, "http_server");
+
+    // SSE transport
+    let sse_config = McpServerConfig {
+        name: "sse_server".to_string(),
+        transport: McpTransport::Sse {
+            url: "http://localhost:8081/sse".to_string(),
+            token: None,
+        },
+        proxy: None,
+        required: false,
+    };
+    assert_eq!(sse_config.name, "sse_server");
+
+    // STDIO transport
+    let stdio_config = McpServerConfig {
+        name: "stdio_server".to_string(),
+        transport: McpTransport::Stdio {
+            command: "mcp-server".to_string(),
+            args: vec!["--port".to_string(), "8082".to_string()],
+            envs: HashMap::new(),
+        },
+        proxy: None,
+        required: false,
+    };
+    assert_eq!(stdio_config.name, "stdio_server");
+}
+
+// Integration Pattern Tests
+
+#[tokio::test]
+async fn test_complete_workflow() {
+    let mock_server = create_mock_server().await;
+
+    // 1. Initialize configuration
+    let config = McpConfig {
+        servers: vec![McpServerConfig {
+            name: "integration_test".to_string(),
+            transport: McpTransport::Streamable {
+                url: mock_server.url(),
+                token: None,
+            },
+            proxy: None,
+            required: false,
+        }],
+        pool: Default::default(),
+        proxy: None,
+        warmup: Vec::new(),
+        inventory: Default::default(),
+    };
+
+    // 2. Connect to server
+    let manager = McpManager::with_defaults(config)
+        .await
+        .expect("Should connect to mock server");
+
+    // 3. Verify server connection
+    let servers = manager.list_servers();
+    assert_eq!(servers.len(), 1);
+    assert_eq!(servers[0], "integration_test");
+
+    // 4. Check available tools
+    let tools = manager.list_tools();
+    assert_eq!(tools.len(), 2);
+
+    // 5. Verify specific tools exist
+    assert!(manager.has_tool("brave_web_search"));
+    assert!(manager.has_tool("brave_local_search"));
+    assert!(!manager.has_tool("nonexistent_tool"));
+
+    // 6. Execute a tool
+    let result = manager
+        .call_tool(
+            "brave_web_search",
+            Some(
+                json!({
+                    "query": "SGLang router MCP integration",
+                    "count": 1
+                })
+                .as_object()
+                .unwrap()
+                .clone(),
+            ),
+        )
+        .await;
+
+    assert!(result.is_ok(), "Tool execution should succeed");
+    let response = result.unwrap();
+    assert!(!response.content.is_empty(), "Should return content");
+
+    // 7. Clean shutdown
+    manager.shutdown().await;
+
+    let capabilities = [
+        "MCP server initialization",
+        "Tool server connection and discovery",
+        "Tool availability checking",
+        "Tool execution",
+        "Error handling and robustness",
+        "Multi-server support",
+        "Schema adaptation",
+        "Mock server integration (no external dependencies)",
+    ];
+
+    assert_eq!(capabilities.len(), 8);
+}

third_party/sglang/sgl-model-gateway/tests/metrics_aggregator_test.rs

@@ -0,0 +1,271 @@
+use smg::core::metrics_aggregator::{aggregate_metrics, MetricPack};
+
+#[test]
+fn test_aggregate_simple() {
+    let pack1 = MetricPack {
+        labels: vec![("source".to_string(), "worker1".to_string())],
+        metrics_text: r#"
+# HELP http_requests_total The total number of HTTP requests.
+# TYPE http_requests_total counter
+http_requests_total{method="post",code="200"} 1027
+http_requests_total{method="post",code="400"} 3
+"#
+        .to_string(),
+    };
+    let pack2 = MetricPack {
+        labels: vec![("source".to_string(), "worker2".to_string())],
+        metrics_text: r#"
+# HELP http_requests_total The total number of HTTP requests.
+# TYPE http_requests_total counter
+http_requests_total{method="post",code="200"} 500
+"#
+        .to_string(),
+    };
+
+    let result = aggregate_metrics(vec![pack1, pack2]).unwrap();
+    let expected = r#"# HELP http_requests_total The total number of HTTP requests.
+# TYPE http_requests_total counter
+http_requests_total{code="200",method="post",source="worker1"} 1027
+http_requests_total{code="400",method="post",source="worker1"} 3
+http_requests_total{code="200",method="post",source="worker2"} 500
+"#;
+    assert_eq!(result.trim(), expected.trim());
+}
+
+#[test]
+fn test_aggregate_multiple_metrics() {
+    let pack1 = MetricPack {
+        labels: vec![("source".to_string(), "w1".to_string())],
+        metrics_text: r#"
+# TYPE metric_a gauge
+metric_a{dim="x"} 1.0
+# TYPE metric_b_total counter
+metric_b_total 10
+"#
+        .to_string(),
+    };
+    let pack2 = MetricPack {
+        labels: vec![("source".to_string(), "w2".to_string())],
+        metrics_text: r#"
+# TYPE metric_a gauge
+metric_a{dim="y"} 2.0
+"#
+        .to_string(),
+    };
+
+    let result = aggregate_metrics(vec![pack1, pack2]).unwrap();
+    let expected = r#"# TYPE metric_a gauge
+metric_a{dim="x",source="w1"} 1
+metric_a{dim="y",source="w2"} 2
+
+# TYPE metric_b_total counter
+metric_b_total{source="w1"} 10
+"#;
+    assert_eq_sorted(&result, expected);
+}
+
+#[test]
+fn test_empty_input() {
+    let result = aggregate_metrics(vec![]).unwrap();
+    assert_eq!(result, "");
+}
+
+#[test]
+fn test_invalid_metrics_are_skipped() {
+    let pack1 = MetricPack {
+        labels: vec![("source".to_string(), "worker1".to_string())],
+        metrics_text: "invalid metrics text".to_string(),
+    };
+    let pack2 = MetricPack {
+        labels: vec![("source".to_string(), "worker2".to_string())],
+        metrics_text: "# TYPE valid_metric gauge\nvalid_metric 123\n".to_string(),
+    };
+    let result = aggregate_metrics(vec![pack1, pack2]).unwrap();
+    let expected = r#"# TYPE valid_metric gauge
+valid_metric{source="worker2"} 123
+"#;
+    assert_eq!(result.trim(), expected.trim());
+}
+
+#[test]
+fn test_real() {
+    let pack1 = MetricPack {
+        labels: vec![("source".to_string(), "worker1".to_string())],
+        // https://docs.sglang.io/references/production_metrics.html
+        metrics_text: r###"# HELP sglang:prompt_tokens_total Number of prefill tokens processed.
+# TYPE sglang:prompt_tokens_total counter
+sglang:prompt_tokens_total{model_name="meta-llama/Llama-3.1-8B-Instruct"} 8.128902e+06
+# HELP sglang:generation_tokens_total Number of generation tokens processed.
+# TYPE sglang:generation_tokens_total counter
+sglang:generation_tokens_total{model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.557572e+06
+# HELP sglang:token_usage The token usage
+# TYPE sglang:token_usage gauge
+sglang:token_usage{model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.28
+# HELP sglang:cache_hit_rate The cache hit rate
+# TYPE sglang:cache_hit_rate gauge
+sglang:cache_hit_rate{model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.007507552643049313
+# HELP sglang:time_to_first_token_seconds Histogram of time to first token in seconds.
+# TYPE sglang:time_to_first_token_seconds histogram
+sglang:time_to_first_token_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct"} 2.3518979474117756e+06
+sglang:time_to_first_token_seconds_bucket{le="0.001",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0
+sglang:time_to_first_token_seconds_bucket{le="0.005",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0
+sglang:time_to_first_token_seconds_bucket{le="0.01",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0
+sglang:time_to_first_token_seconds_bucket{le="+Inf",model_name="meta-llama/Llama-3.1-8B-Instruct"} 11008.0
+sglang:time_to_first_token_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct"} 11008.0
+# HELP sglang:e2e_request_latency_seconds Histogram of End-to-end request latency in seconds
+# TYPE sglang:e2e_request_latency_seconds histogram
+sglang:e2e_request_latency_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct"} 3.116093850019932e+06
+sglang:e2e_request_latency_seconds_bucket{le="0.3",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0
+sglang:e2e_request_latency_seconds_bucket{le="0.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0
+sglang:e2e_request_latency_seconds_bucket{le="0.8",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0
+sglang:e2e_request_latency_seconds_bucket{le="+Inf",model_name="meta-llama/Llama-3.1-8B-Instruct"} 11228.0
+sglang:e2e_request_latency_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct"} 11228.0
+# HELP sglang:time_per_output_token_seconds Histogram of time per output token in seconds.
+# TYPE sglang:time_per_output_token_seconds histogram
+sglang:time_per_output_token_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct"} 866964.5791549598
+sglang:time_per_output_token_seconds_bucket{le="0.005",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1.0
+sglang:time_per_output_token_seconds_bucket{le="0.01",model_name="meta-llama/Llama-3.1-8B-Instruct"} 73.0
+sglang:time_per_output_token_seconds_bucket{le="0.015",model_name="meta-llama/Llama-3.1-8B-Instruct"} 382.0
+sglang:time_per_output_token_seconds_bucket{le="+Inf",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.400757e+06
+sglang:time_per_output_token_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.400757e+06
+# HELP sglang:func_latency_seconds Function latency in seconds
+# TYPE sglang:func_latency_seconds histogram
+sglang:func_latency_seconds_sum{name="generate_request"} 4.514771912145079
+sglang:func_latency_seconds_bucket{le="0.05",name="generate_request"} 14006.0
+sglang:func_latency_seconds_bucket{le="0.07500000000000001",name="generate_request"} 14006.0
+sglang:func_latency_seconds_bucket{le="0.1125",name="generate_request"} 14006.0
+sglang:func_latency_seconds_bucket{le="0.16875",name="generate_request"} 14006.0
+sglang:func_latency_seconds_bucket{le="+Inf",name="generate_request"} 14007.0
+sglang:func_latency_seconds_count{name="generate_request"} 14007.0
+# HELP sglang:num_running_reqs The number of running requests
+# TYPE sglang:num_running_reqs gauge
+sglang:num_running_reqs{model_name="meta-llama/Llama-3.1-8B-Instruct"} 162.0
+# HELP sglang:num_used_tokens The number of used tokens
+# TYPE sglang:num_used_tokens gauge
+sglang:num_used_tokens{model_name="meta-llama/Llama-3.1-8B-Instruct"} 123859.0
+# HELP sglang:gen_throughput The generate throughput (token/s)
+# TYPE sglang:gen_throughput gauge
+sglang:gen_throughput{model_name="meta-llama/Llama-3.1-8B-Instruct"} 86.50814177726902
+# HELP sglang:num_queue_reqs The number of requests in the waiting queue
+# TYPE sglang:num_queue_reqs gauge
+sglang:num_queue_reqs{model_name="meta-llama/Llama-3.1-8B-Instruct"} 2826.0
+"###.to_string(),
+    };
+    let pack2 = MetricPack {
+        labels: vec![("source".to_string(), "worker2".to_string())],
+        metrics_text: pack1.metrics_text.clone(),
+    };
+    let result = aggregate_metrics(vec![pack1, pack2]).unwrap();
+    let expected = r###"# HELP sglang_token_usage The token usage
+# TYPE sglang_token_usage gauge
+sglang_token_usage{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1"} 0.28
+sglang_token_usage{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2"} 0.28
+
+# HELP sglang_time_to_first_token_seconds Histogram of time to first token in seconds.
+# TYPE sglang_time_to_first_token_seconds histogram
+sglang_time_to_first_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1",le="0.001"} 0
+sglang_time_to_first_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1",le="0.005"} 0
+sglang_time_to_first_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1",le="0.01"} 0
+sglang_time_to_first_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1",le="+Inf"} 11008
+sglang_time_to_first_token_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1"} 2351897.9474117756
+sglang_time_to_first_token_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1"} 11008
+sglang_time_to_first_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2",le="0.001"} 0
+sglang_time_to_first_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2",le="0.005"} 0
+sglang_time_to_first_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2",le="0.01"} 0
+sglang_time_to_first_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2",le="+Inf"} 11008
+sglang_time_to_first_token_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2"} 2351897.9474117756
+sglang_time_to_first_token_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2"} 11008
+
+# HELP sglang_time_per_output_token_seconds Histogram of time per output token in seconds.
+# TYPE sglang_time_per_output_token_seconds histogram
+sglang_time_per_output_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1",le="0.005"} 1
+sglang_time_per_output_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1",le="0.01"} 73
+sglang_time_per_output_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1",le="0.015"} 382
+sglang_time_per_output_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1",le="+Inf"} 7400757
+sglang_time_per_output_token_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1"} 866964.5791549598
+sglang_time_per_output_token_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1"} 7400757
+sglang_time_per_output_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2",le="0.005"} 1
+sglang_time_per_output_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2",le="0.01"} 73
+sglang_time_per_output_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2",le="0.015"} 382
+sglang_time_per_output_token_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2",le="+Inf"} 7400757
+sglang_time_per_output_token_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2"} 866964.5791549598
+sglang_time_per_output_token_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2"} 7400757
+
+# HELP sglang_func_latency_seconds Function latency in seconds
+# TYPE sglang_func_latency_seconds histogram
+sglang_func_latency_seconds_bucket{name="generate_request",source="worker1",le="0.05"} 14006
+sglang_func_latency_seconds_bucket{name="generate_request",source="worker1",le="0.07500000000000001"} 14006
+sglang_func_latency_seconds_bucket{name="generate_request",source="worker1",le="0.1125"} 14006
+sglang_func_latency_seconds_bucket{name="generate_request",source="worker1",le="0.16875"} 14006
+sglang_func_latency_seconds_bucket{name="generate_request",source="worker1",le="+Inf"} 14007
+sglang_func_latency_seconds_sum{name="generate_request",source="worker1"} 4.514771912145079
+sglang_func_latency_seconds_count{name="generate_request",source="worker1"} 14007
+sglang_func_latency_seconds_bucket{name="generate_request",source="worker2",le="0.05"} 14006
+sglang_func_latency_seconds_bucket{name="generate_request",source="worker2",le="0.07500000000000001"} 14006
+sglang_func_latency_seconds_bucket{name="generate_request",source="worker2",le="0.1125"} 14006
+sglang_func_latency_seconds_bucket{name="generate_request",source="worker2",le="0.16875"} 14006
+sglang_func_latency_seconds_bucket{name="generate_request",source="worker2",le="+Inf"} 14007
+sglang_func_latency_seconds_sum{name="generate_request",source="worker2"} 4.514771912145079
+sglang_func_latency_seconds_count{name="generate_request",source="worker2"} 14007
+
+# HELP sglang_num_used_tokens The number of used tokens
+# TYPE sglang_num_used_tokens gauge
+sglang_num_used_tokens{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1"} 123859
+sglang_num_used_tokens{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2"} 123859
+
+# HELP sglang_cache_hit_rate The cache hit rate
+# TYPE sglang_cache_hit_rate gauge
+sglang_cache_hit_rate{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1"} 0.007507552643049313
+sglang_cache_hit_rate{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2"} 0.007507552643049313
+
+# HELP sglang_num_queue_reqs The number of requests in the waiting queue
+# TYPE sglang_num_queue_reqs gauge
+sglang_num_queue_reqs{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1"} 2826
+sglang_num_queue_reqs{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2"} 2826
+
+# HELP sglang_generation_tokens_total Number of generation tokens processed.
+# TYPE sglang_generation_tokens_total counter
+sglang_generation_tokens_total{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1"} 7557572
+sglang_generation_tokens_total{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2"} 7557572
+
+# HELP sglang_num_running_reqs The number of running requests
+# TYPE sglang_num_running_reqs gauge
+sglang_num_running_reqs{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1"} 162
+sglang_num_running_reqs{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2"} 162
+
+# HELP sglang_e2e_request_latency_seconds Histogram of End-to-end request latency in seconds
+# TYPE sglang_e2e_request_latency_seconds histogram
+sglang_e2e_request_latency_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1",le="0.3"} 0
+sglang_e2e_request_latency_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1",le="0.5"} 6
+sglang_e2e_request_latency_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1",le="0.8"} 6
+sglang_e2e_request_latency_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1",le="+Inf"} 11228
+sglang_e2e_request_latency_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1"} 3116093.850019932
+sglang_e2e_request_latency_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1"} 11228
+sglang_e2e_request_latency_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2",le="0.3"} 0
+sglang_e2e_request_latency_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2",le="0.5"} 6
+sglang_e2e_request_latency_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2",le="0.8"} 6
+sglang_e2e_request_latency_seconds_bucket{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2",le="+Inf"} 11228
+sglang_e2e_request_latency_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2"} 3116093.850019932
+sglang_e2e_request_latency_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2"} 11228
+
+# HELP sglang_gen_throughput The generate throughput (token/s)
+# TYPE sglang_gen_throughput gauge
+sglang_gen_throughput{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1"} 86.50814177726902
+sglang_gen_throughput{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2"} 86.50814177726902
+
+# HELP sglang_prompt_tokens_total Number of prefill tokens processed.
+# TYPE sglang_prompt_tokens_total counter
+sglang_prompt_tokens_total{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker1"} 8128902
+sglang_prompt_tokens_total{model_name="meta-llama/Llama-3.1-8B-Instruct",source="worker2"} 8128902"###;
+    println!("result=\n{result}");
+    assert_eq_sorted(result.trim(), expected.trim());
+}
+
+fn assert_eq_sorted(result: &str, expected: &str) {
+    // Split into lines and sort to handle BTreeMap ordering issues between test environments
+    let mut result_lines: Vec<_> = result.trim().lines().map(|l| l.trim()).collect();
+    let mut expected_lines: Vec<_> = expected.trim().lines().map(|l| l.trim()).collect();
+    result_lines.sort();
+    expected_lines.sort();
+    assert_eq!(result_lines, expected_lines);
+}

third_party/sglang/sgl-model-gateway/tests/otel_tracing_test.rs

@@ -0,0 +1,380 @@
+mod common;
+
+use std::{
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+    time::Duration,
+};
+
+use axum::{body::Body, extract::Request, http::StatusCode};
+use common::mock_worker::{HealthStatus, MockWorker, MockWorkerConfig, WorkerType};
+use opentelemetry_proto::tonic::collector::trace::v1::{
+    trace_service_server::{TraceService, TraceServiceServer},
+    ExportTraceServiceRequest, ExportTraceServiceResponse,
+};
+use portpicker::pick_unused_port;
+use serde_json::json;
+use serial_test::serial;
+use smg::{
+    config::{RouterConfig, TraceConfig},
+    core::Job,
+    observability::{logging, otel_trace},
+    routers::RouterFactory,
+};
+use tokio::sync::oneshot;
+use tonic::metadata::MetadataMap;
+use tonic_v12::{transport::Server, Request as TonicRequest, Response, Status};
+use tower::ServiceExt;
+use tracing::info_span;
+use tracing_subscriber::prelude::*;
+
+#[derive(Clone)]
+struct TestOtelCollector {
+    span_count: Arc<AtomicUsize>,
+}
+
+impl TestOtelCollector {
+    fn new() -> Self {
+        Self {
+            span_count: Arc::new(AtomicUsize::new(0)),
+        }
+    }
+
+    fn get_span_count(&self) -> usize {
+        self.span_count.load(Ordering::SeqCst)
+    }
+}
+
+#[tonic_v12::async_trait]
+impl TraceService for TestOtelCollector {
+    async fn export(
+        &self,
+        request: TonicRequest<ExportTraceServiceRequest>,
+    ) -> Result<Response<ExportTraceServiceResponse>, Status> {
+        let req = request.into_inner();
+
+        let mut total_spans = 0;
+
+        for resource_span in &req.resource_spans {
+            for scope_span in &resource_span.scope_spans {
+                total_spans += scope_span.spans.len();
+            }
+        }
+
+        self.span_count.fetch_add(total_spans, Ordering::SeqCst);
+
+        Ok(Response::new(ExportTraceServiceResponse {
+            partial_success: None,
+        }))
+    }
+}
+
+async fn start_collector(
+    port: u16,
+    shutdown_rx: oneshot::Receiver<()>,
+) -> Result<TestOtelCollector, Box<dyn std::error::Error>> {
+    let addr = format!("0.0.0.0:{}", port).parse()?;
+    let collector = TestOtelCollector::new();
+    let collector_clone = collector.clone();
+
+    tokio::spawn(async move {
+        let _ = Server::builder()
+            .add_service(TraceServiceServer::new(collector_clone))
+            .serve_with_shutdown(addr, async {
+                shutdown_rx.await.ok();
+            })
+            .await;
+    });
+
+    tokio::time::sleep(Duration::from_millis(200)).await;
+
+    Ok(collector)
+}
+
+#[tokio::test]
+#[serial]
+async fn test_router_with_tracing() {
+    // 1. Start the OTLP collector
+    let port = pick_unused_port().expect("Failed to pick unused port");
+    let (shutdown_tx, shutdown_rx) = oneshot::channel();
+    let collector = start_collector(port, shutdown_rx)
+        .await
+        .expect("Failed to start collector");
+    let collector_endpoint = format!("0.0.0.0:{}", port);
+    println!("OTLP Collector started on: {}", collector_endpoint);
+
+    // 2. create the mock worker
+    let mut mock_worker = MockWorker::new(MockWorkerConfig {
+        port: 0,
+        worker_type: WorkerType::Regular,
+        health_status: HealthStatus::Healthy,
+        response_delay_ms: 0,
+        fail_rate: 0.0,
+    });
+
+    let worker_url = mock_worker.start().await.unwrap();
+    tokio::time::sleep(Duration::from_millis(200)).await;
+    println!("Mock worker started on: {}", worker_url);
+
+    // 3. create router config and enable tracing
+    let router_config = RouterConfig::builder()
+        .regular_mode(vec![worker_url.clone()])
+        .random_policy()
+        .host("0.0.0.0")
+        .port(0)
+        .max_payload_size(256 * 1024 * 1024)
+        .request_timeout_secs(60)
+        .worker_startup_timeout_secs(1)
+        .worker_startup_check_interval_secs(1)
+        .max_concurrent_requests(64)
+        .queue_timeout_secs(60)
+        .enable_trace(&collector_endpoint)
+        .build_unchecked();
+
+    // 4. Initialize the OTLP client (check if already initialized by another test)
+    let otel_initialized_by_this_test = if !otel_trace::is_otel_enabled() {
+        let init_result = otel_trace::otel_tracing_init(true, Some(&collector_endpoint));
+        assert!(
+            init_result.is_ok(),
+            "Failed to initialize OTEL: {:?}",
+            init_result.err()
+        );
+        println!("OpenTelemetry initialized successfully");
+        true
+    } else {
+        println!(
+            "OpenTelemetry already initialized by previous test (spans will go to that collector)"
+        );
+        false
+    };
+
+    let trace_config = TraceConfig {
+        enable_trace: true,
+        otlp_traces_endpoint: collector_endpoint.clone(),
+    };
+    let _log_guard = logging::init_logging(
+        logging::LoggingConfig {
+            level: tracing::Level::INFO,
+            json_format: false,
+            log_dir: None,
+            colorize: false,
+            log_file_name: "test-otel".to_string(),
+            log_targets: Some(vec!["smg".to_string()]),
+        },
+        Some(trace_config),
+    );
+    println!("Logging initialized with OTEL layer");
+
+    // 5. Create a span and sleep for a while
+    let _span = info_span!(target: "smg::otel-trace", "test_router_with_tracing");
+    tokio::time::sleep(Duration::from_secs(1)).await;
+    drop(_span);
+
+    // 6. create app context and router
+    let app_context = common::create_test_context(router_config.clone()).await;
+
+    // 7. initialize worker
+    let job_queue = app_context
+        .worker_job_queue
+        .get()
+        .expect("JobQueue should be initialized");
+
+    let job = Job::InitializeWorkersFromConfig {
+        router_config: Box::new(router_config.clone()),
+    };
+
+    job_queue
+        .submit(job)
+        .await
+        .expect("Failed to submit worker init job");
+
+    // 8. wait for worker initialization
+    tokio::time::sleep(Duration::from_millis(1000)).await;
+    println!("Workers initialized");
+
+    // 9. create router
+    let router = RouterFactory::create_router(&app_context)
+        .await
+        .expect("Failed to create router");
+
+    println!("Router created");
+
+    // 10. create app (middleware::create_logging_layer() will use the already initialized OTEL layer)
+    let app =
+        common::test_app::create_test_app_with_context(Arc::from(router), app_context.clone());
+
+    println!("App created with logging middleware");
+
+    // 10. send request
+    let request_body = json!({
+        "model": "test-model",
+        "messages": [
+            {"role": "user", "content": "Hello, test OpenTelemetry tracing!"}
+        ],
+        "temperature": 0.7,
+        "max_tokens": 50
+    });
+
+    let request = Request::builder()
+        .method("POST")
+        .uri("/v1/chat/completions")
+        .header("content-type", "application/json")
+        .body(Body::from(request_body.to_string()))
+        .unwrap();
+
+    println!("Sending request to router...");
+    let response = app.oneshot(request).await.unwrap();
+
+    assert_eq!(response.status(), StatusCode::OK, "Request should succeed");
+
+    println!("Request completed successfully");
+    drop(response);
+
+    // 11. Wait for spans to be exported
+    match otel_trace::flush_spans_async().await {
+        Ok(_) => println!("Spans flushed successfully"),
+        Err(e) => println!("Failed to flush spans: {:?}", e),
+    }
+
+    // 12. Verify that the spans were exported to the collector
+    let span_count = collector.get_span_count();
+    println!("Total spans received by collector: {}", span_count);
+
+    // Only assert span count if we initialized OTEL with our own collector
+    // When OTEL was pre-initialized by another test, spans go to that collector instead
+    if otel_initialized_by_this_test {
+        assert!(
+            span_count == 2,
+            "Expected to receive at least 2 span, but got {}. \
+            This indicates that tracing data is not being exported to the OTLP collector.",
+            span_count
+        );
+        println!("Test passed! Collector received {} spans", span_count);
+    } else {
+        println!(
+            "Skipping span count assertion - OTEL was pre-initialized by another test. \
+            Spans went to that collector. Received {} spans on this test's collector.",
+            span_count
+        );
+    }
+
+    // 13. cleanup
+    let _ = shutdown_tx.send(());
+    mock_worker.stop().await;
+
+    println!("Cleanup completed");
+}
+
+// ============================================================================
+// gRPC Trace Context Injection Tests
+// ============================================================================
+
+/// Comprehensive test for gRPC trace context injection.
+///
+/// This test validates:
+/// 1. W3C trace context headers are properly injected into gRPC metadata
+/// 2. traceparent format is correct (version-traceid-spanid-flags)
+/// 3. All metadata keys are lowercase (gRPC requirement)
+///
+/// Note: This test handles the case where OTEL may already be initialized
+/// by a previous test (since tests run sequentially with #[serial]).
+#[tokio::test]
+#[serial]
+async fn test_grpc_trace_context_injection() {
+    // 1. Start the OTLP collector (needed even if OTEL is already initialized,
+    //    as a target for any spans that might be exported)
+    let port = pick_unused_port().expect("Failed to pick unused port");
+    let (shutdown_tx, shutdown_rx) = oneshot::channel();
+    let _collector = start_collector(port, shutdown_rx)
+        .await
+        .expect("Failed to start collector");
+    let collector_endpoint = format!("0.0.0.0:{}", port);
+
+    // 2. Initialize OTEL if not already enabled
+    // Note: otel_tracing_init will fail if already initialized (OnceLock),
+    // but that's fine - we just need OTEL to be enabled
+    let already_enabled = otel_trace::is_otel_enabled();
+    if !already_enabled {
+        let init_result = otel_trace::otel_tracing_init(true, Some(&collector_endpoint));
+        assert!(
+            init_result.is_ok(),
+            "Failed to initialize OTEL: {:?}",
+            init_result.err()
+        );
+    }
+
+    // Verify OTEL is enabled (either from this test or a previous one)
+    assert!(otel_trace::is_otel_enabled(), "OTEL should be enabled");
+
+    // 3. Set up tracing subscriber with OTEL layer
+    let otel_layer = otel_trace::get_otel_layer().expect("Failed to get OTEL layer");
+    let subscriber = tracing_subscriber::registry().with(otel_layer);
+
+    // 4. Test within a span context
+    tracing::subscriber::with_default(subscriber, || {
+        // Create a span that will be exported to OTEL
+        let span = info_span!(target: "smg::otel-trace", "test_grpc_span");
+        let _guard = span.enter();
+
+        // Create empty gRPC metadata
+        let mut metadata = MetadataMap::new();
+
+        // Inject trace context
+        otel_trace::inject_trace_context_grpc(&mut metadata);
+
+        // === Test 1: Verify traceparent header was injected ===
+        let traceparent = metadata.get("traceparent");
+        assert!(
+            traceparent.is_some(),
+            "traceparent header should be present in gRPC metadata"
+        );
+
+        // === Test 2: Verify traceparent format (version-traceid-spanid-flags) ===
+        let traceparent_value = traceparent.unwrap().to_str().unwrap();
+        let parts: Vec<&str> = traceparent_value.split('-').collect();
+        assert_eq!(
+            parts.len(),
+            4,
+            "traceparent should have 4 parts: version-traceid-spanid-flags"
+        );
+        assert_eq!(parts[0], "00", "traceparent version should be 00");
+        assert_eq!(parts[1].len(), 32, "trace ID should be 32 hex characters");
+        assert_eq!(parts[2].len(), 16, "span ID should be 16 hex characters");
+
+        println!("Successfully injected traceparent: {}", traceparent_value);
+
+        // === Test 3: Verify all keys are lowercase (gRPC metadata requirement) ===
+        for key_and_value in metadata.iter() {
+            match key_and_value {
+                tonic::metadata::KeyAndValueRef::Ascii(key, _) => {
+                    let key_str = key.as_str();
+                    assert_eq!(
+                        key_str,
+                        key_str.to_lowercase(),
+                        "gRPC metadata key '{}' should be lowercase",
+                        key_str
+                    );
+                }
+                tonic::metadata::KeyAndValueRef::Binary(key, _) => {
+                    let key_str = key.as_str();
+                    assert_eq!(
+                        key_str,
+                        key_str.to_lowercase(),
+                        "gRPC metadata key '{}' should be lowercase",
+                        key_str
+                    );
+                }
+            }
+        }
+
+        println!("All gRPC metadata keys are lowercase as required");
+    });
+
+    // Cleanup - don't shutdown OTEL since tests share global state (OnceLock)
+    // and other tests may need to use the already-initialized OTEL
+    let _ = shutdown_tx.send(());
+
+    println!("test_grpc_trace_context_injection: All assertions passed!");
+}

third_party/sglang/sgl-model-gateway/tests/reliability/circuit_breaker_test.rs

@@ -0,0 +1,243 @@
+//! Circuit breaker integration tests
+//!
+//! Tests for circuit breaker behavior: opening, half-open state, recovery, and isolation.
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use smg::config::{CircuitBreakerConfig, RetryConfig, RouterConfig};
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+#[cfg(test)]
+mod circuit_breaker_tests {
+    use super::*;
+
+    /// Test that circuit breaker opens after consecutive failures
+    #[tokio::test]
+    async fn test_circuit_breaker_opens_after_failures() {
+        let config = TestRouterConfig::round_robin_with_circuit_breaker(
+            3200,
+            CircuitBreakerConfig {
+                failure_threshold: 3,
+                success_threshold: 2,
+                timeout_duration_secs: 2,
+                window_duration_secs: 10,
+            },
+        );
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![TestWorkerConfig::flaky(19100, 1.0)], // Always fail
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Make requests until we see failures (500) and then circuit breaker opens (503)
+        let mut saw_500 = false;
+        let mut saw_503 = false;
+
+        for _ in 0..10 {
+            let payload = json!({
+                "text": "Test circuit breaker",
+                "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();
+            match resp.status() {
+                StatusCode::INTERNAL_SERVER_ERROR => saw_500 = true,
+                StatusCode::SERVICE_UNAVAILABLE => saw_503 = true,
+                _ => {}
+            }
+        }
+
+        // Should see 500 (worker error) and eventually 503 (circuit breaker open)
+        assert!(
+            saw_500 || saw_503,
+            "Should see either 500 (worker error) or 503 (circuit breaker open)"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test circuit breaker with disabled flag
+    #[tokio::test]
+    async fn test_circuit_breaker_disabled() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(3201)
+            .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)
+            .disable_circuit_breaker()
+            .disable_retries()
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![TestWorkerConfig::flaky(19101, 1.0)], // Always fail
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // With circuit breaker disabled, should always see 500 (never 503)
+        for _ in 0..5 {
+            let payload = json!({
+                "text": "Test disabled CB",
+                "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();
+            // With CB disabled, we expect 500 errors (not 503 from CB)
+            assert_eq!(
+                resp.status(),
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Should get 500 with CB disabled, not CB-related 503"
+            );
+        }
+        ctx.shutdown().await;
+    }
+
+    /// Test circuit breaker per-worker isolation
+    #[tokio::test]
+    async fn test_circuit_breaker_per_worker_isolation() {
+        let config = TestRouterConfig::round_robin_with_circuit_breaker(
+            3202,
+            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(19102, 1.0), // Always fail
+                TestWorkerConfig::healthy(19103),    // Always succeed
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Send enough requests to trigger CB on the failing worker
+        // The healthy worker should continue to serve requests
+        let mut success_count = 0;
+        for _ in 0..20 {
+            let payload = json!({
+                "text": "Test isolation",
+                "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;
+            }
+        }
+
+        // After CB opens on failing worker, healthy worker should handle all requests
+        // We should see at least some successes
+        assert!(
+            success_count > 0,
+            "Should have some successful requests from healthy worker"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test circuit breaker with retries enabled
+    #[tokio::test]
+    async fn test_circuit_breaker_with_retries() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(3203)
+            .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()
+            })
+            .circuit_breaker_config(CircuitBreakerConfig {
+                failure_threshold: 2,
+                success_threshold: 1,
+                timeout_duration_secs: 2,
+                window_duration_secs: 10,
+            })
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(19104, 1.0), // Always fail
+                TestWorkerConfig::healthy(19105),    // Always succeed
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // With retries, requests should succeed by retrying on 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;
+    }
+}

third_party/sglang/sgl-model-gateway/tests/reliability/fault_tolerance_test.rs

@@ -0,0 +1,309 @@
+//! Fault tolerance integration tests
+//!
+//! Tests for system resilience: worker failures, network issues, and recovery.
+
+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::{CircuitBreakerConfig, RetryConfig};
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+#[cfg(test)]
+mod fault_tolerance_tests {
+    use super::*;
+
+    /// Test that requests are rerouted when a worker fails
+    #[tokio::test]
+    async fn test_worker_failure_reroute() {
+        let config = TestRouterConfig::round_robin_with_reliability(
+            4100,
+            RetryConfig {
+                max_retries: 3,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 100,
+                ..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(20100, 1.0), // Always fails
+                TestWorkerConfig::healthy(20101),    // 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!("Fault tolerance test {}", 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 reroute to healthy worker"
+            );
+        }
+
+        ctx.shutdown().await;
+    }
+
+    /// Test behavior when all workers are temporarily unavailable
+    #[tokio::test]
+    async fn test_all_workers_temporarily_failing() {
+        let config = TestRouterConfig::round_robin_with_retry(
+            4101,
+            RetryConfig {
+                max_retries: 2,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 50,
+                ..Default::default()
+            },
+        );
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(20102, 1.0), // Always fails
+                TestWorkerConfig::flaky(20103, 1.0), // Always fails
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        let payload = json!({
+            "text": "Test with all failing 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 fail when all workers are failing
+        assert!(
+            resp.status() == StatusCode::INTERNAL_SERVER_ERROR
+                || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
+            "Request should fail when all workers are failing, got {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test graceful handling of slow workers
+    #[tokio::test]
+    async fn test_slow_worker_handling() {
+        let config = TestRouterConfig::round_robin(4102);
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::slow(20104, 500), // Slow worker
+                TestWorkerConfig::healthy(20105),   // Fast worker
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Send concurrent requests
+        let mut handles = Vec::new();
+        let success_count = Arc::new(AtomicUsize::new(0));
+
+        for i in 0..10 {
+            let app_clone = app.clone();
+            let success_clone = Arc::clone(&success_count);
+
+            let handle = tokio::spawn(async move {
+                let payload = json!({
+                    "text": format!("Slow worker test {}", 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 eventually succeed
+        assert_eq!(
+            success_count.load(Ordering::SeqCst),
+            10,
+            "All requests should succeed despite slow worker"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test circuit breaker prevents cascading failures
+    #[tokio::test]
+    async fn test_circuit_breaker_prevents_cascade() {
+        let config = TestRouterConfig::round_robin_with_reliability(
+            4103,
+            RetryConfig {
+                max_retries: 3,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 50,
+                ..Default::default()
+            },
+            CircuitBreakerConfig {
+                failure_threshold: 2,
+                success_threshold: 1,
+                timeout_duration_secs: 5,
+                window_duration_secs: 10,
+            },
+        );
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(20106, 1.0), // Failing worker
+                TestWorkerConfig::healthy(20107),    // Healthy worker
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+        let mut success_count = 0;
+
+        // Send many requests - after CB opens, all should route to healthy worker
+        for i in 0..20 {
+            let payload = json!({
+                "text": format!("Circuit breaker cascade test {}", 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;
+            }
+        }
+
+        // Most requests should succeed after CB opens on failing worker
+        assert!(
+            success_count >= 15,
+            "Most requests should succeed after circuit breaker opens, got {} successes",
+            success_count
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test recovery after worker comes back online (simulated via healthy worker)
+    #[tokio::test]
+    async fn test_system_stability_under_partial_failure() {
+        let config = TestRouterConfig::round_robin_with_reliability(
+            4104,
+            RetryConfig {
+                max_retries: 2,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 50,
+                ..Default::default()
+            },
+            CircuitBreakerConfig {
+                failure_threshold: 3,
+                success_threshold: 1,
+                timeout_duration_secs: 2,
+                window_duration_secs: 10,
+            },
+        );
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(20108, 0.5), // 50% failure rate
+                TestWorkerConfig::healthy(20109),    // Always succeeds
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+        let mut success_count = 0;
+
+        // System should maintain stability with partial failures
+        for i in 0..30 {
+            let payload = json!({
+                "text": format!("Stability test {}", 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;
+            }
+        }
+
+        // With retries and one healthy worker, most should succeed
+        assert!(
+            success_count >= 25,
+            "System should maintain stability under partial failure, got {} successes",
+            success_count
+        );
+
+        ctx.shutdown().await;
+    }
+}

third_party/sglang/sgl-model-gateway/tests/reliability/mod.rs

@@ -0,0 +1,6 @@
+//! Reliability integration tests
+
+pub mod circuit_breaker_test;
+pub mod fault_tolerance_test;
+pub mod rate_limiting_test;
+pub mod retries_test;

third_party/sglang/sgl-model-gateway/tests/reliability/rate_limiting_test.rs

@@ -0,0 +1,263 @@
+//! Rate limiting integration tests
+//!
+//! Tests for rate limiting and concurrency control.
+
+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 rate_limiting_tests {
+    use super::*;
+
+    /// Test that concurrent requests are handled within limits
+    #[tokio::test]
+    async fn test_concurrent_requests_within_limit() {
+        let config = TestRouterConfig::with_concurrency(3400, 10);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(19300, 50)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Send 5 concurrent requests (within limit of 10)
+        let mut handles = Vec::new();
+        let success_count = Arc::new(AtomicUsize::new(0));
+
+        for i in 0..5 {
+            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);
+        }
+
+        // Wait for all requests to complete
+        for handle in handles {
+            handle.await.unwrap();
+        }
+
+        // All requests should succeed within the limit
+        assert_eq!(
+            success_count.load(Ordering::SeqCst),
+            5,
+            "All concurrent requests should succeed within limit"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test rate limit tokens per second
+    #[tokio::test]
+    async fn test_rate_limit_tokens() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .random_policy()
+            .host("127.0.0.1")
+            .port(3401)
+            .max_payload_size(256 * 1024 * 1024)
+            .request_timeout_secs(600)
+            .worker_startup_timeout_secs(5)
+            .worker_startup_check_interval_secs(1)
+            .max_concurrent_requests(100)
+            .rate_limit_tokens_per_second(50) // 50 tokens/sec
+            .queue_timeout_secs(60)
+            .build_unchecked();
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(19301)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Send requests and verify they succeed
+        let mut success_count = 0;
+        for i in 0..10 {
+            let payload = json!({
+                "text": format!("Rate limited 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;
+            }
+        }
+
+        // Most requests should succeed (some might be rate limited)
+        assert!(
+            success_count >= 5,
+            "At least half of requests should succeed, got {}",
+            success_count
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test unlimited concurrent requests when set to 0
+    #[tokio::test]
+    async fn test_unlimited_concurrent_requests() {
+        let config = TestRouterConfig::with_concurrency(3402, 0); // Unlimited
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(19302, 10)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Send many concurrent requests
+        let mut handles = Vec::new();
+        let success_count = Arc::new(AtomicUsize::new(0));
+
+        for i in 0..20 {
+            let app_clone = app.clone();
+            let success_clone = Arc::clone(&success_count);
+
+            let handle = tokio::spawn(async move {
+                let payload = json!({
+                    "text": format!("Unlimited 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 with unlimited concurrency
+        assert_eq!(
+            success_count.load(Ordering::SeqCst),
+            20,
+            "All requests should succeed with unlimited concurrency"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test queue behavior when requests exceed capacity
+    #[tokio::test]
+    async fn test_queue_behavior() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .random_policy()
+            .host("127.0.0.1")
+            .port(3403)
+            .max_payload_size(256 * 1024 * 1024)
+            .request_timeout_secs(600)
+            .worker_startup_timeout_secs(5)
+            .worker_startup_check_interval_secs(1)
+            .max_concurrent_requests(2) // Very low limit
+            .queue_size(5) // Small queue
+            .queue_timeout_secs(1) // Short timeout
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![TestWorkerConfig::slow(19303, 200)], // Longer delay to test queuing
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Send requests that will exceed capacity
+        let mut handles = Vec::new();
+        let success_count = Arc::new(AtomicUsize::new(0));
+        let error_count = Arc::new(AtomicUsize::new(0));
+
+        for i in 0..10 {
+            let app_clone = app.clone();
+            let success_clone = Arc::clone(&success_count);
+            let error_clone = Arc::clone(&error_count);
+
+            let handle = tokio::spawn(async move {
+                let payload = json!({
+                    "text": format!("Queue test {}", 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);
+                } else {
+                    error_clone.fetch_add(1, Ordering::SeqCst);
+                }
+            });
+
+            handles.push(handle);
+        }
+
+        for handle in handles {
+            handle.await.unwrap();
+        }
+
+        // Some requests should succeed, some might timeout or be rejected
+        let successes = success_count.load(Ordering::SeqCst);
+        let errors = error_count.load(Ordering::SeqCst);
+
+        assert!(
+            successes > 0,
+            "At least some requests should succeed, got {} successes and {} errors",
+            successes,
+            errors
+        );
+
+        ctx.shutdown().await;
+    }
+}

third_party/sglang/sgl-model-gateway/tests/reliability/retries_test.rs

@@ -0,0 +1,214 @@
+//! Retry mechanism integration tests
+//!
+//! Tests for retry behavior: exponential backoff, max retries, and retry-on-failure scenarios.
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use smg::config::{RetryConfig, RouterConfig};
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+#[cfg(test)]
+mod retry_tests {
+    use super::*;
+
+    /// Test that retries succeed when at least one worker is healthy
+    #[tokio::test]
+    async fn test_retry_succeeds_with_healthy_fallback() {
+        let retry_config = RetryConfig {
+            max_retries: 3,
+            initial_backoff_ms: 10,
+            max_backoff_ms: 100,
+            ..Default::default()
+        };
+        let config = TestRouterConfig::round_robin_with_retry(3300, retry_config);
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(19200, 1.0), // First worker always fails
+                TestWorkerConfig::healthy(19201),    // Second worker always succeeds
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Request should succeed via retry to healthy worker
+        let payload = json!({
+            "text": "Test retry to healthy 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"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test that retries are disabled when configured
+    #[tokio::test]
+    async fn test_retries_disabled() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(3301)
+            .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)
+            .disable_retries()
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![TestWorkerConfig::flaky(19202, 1.0)], // Always fail
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // With retries disabled, request should fail immediately
+        let payload = json!({
+            "text": "Test no 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::INTERNAL_SERVER_ERROR,
+            "Request should fail without retries"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test max retries limit
+    #[tokio::test]
+    async fn test_max_retries_limit() {
+        let retry_config = RetryConfig {
+            max_retries: 2,
+            initial_backoff_ms: 10,
+            max_backoff_ms: 50,
+            ..Default::default()
+        };
+        let config = TestRouterConfig::round_robin_with_retry(3302, retry_config);
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![TestWorkerConfig::flaky(19203, 1.0)], // Always fail
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // All retries will fail, should return error after exhausting retries
+        let payload = json!({
+            "text": "Test max retries",
+            "stream": false
+        });
+
+        let start = std::time::Instant::now();
+        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();
+
+        // Should eventually fail after retries
+        assert!(
+            resp.status() == StatusCode::INTERNAL_SERVER_ERROR
+                || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
+            "Should fail after exhausting retries, got {}",
+            resp.status()
+        );
+
+        // Should take some time due to backoff (at least initial_backoff_ms)
+        // With 2 retries and 10ms initial backoff, should take at least 10ms
+        // But don't make this too strict as timing can vary
+        assert!(
+            elapsed.as_millis() >= 5,
+            "Should have some backoff delay, got {}ms",
+            elapsed.as_millis()
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test retry with multiple workers - should eventually find healthy one
+    #[tokio::test]
+    async fn test_retry_finds_healthy_worker() {
+        let retry_config = RetryConfig {
+            max_retries: 5,
+            initial_backoff_ms: 5,
+            max_backoff_ms: 50,
+            ..Default::default()
+        };
+        let config = TestRouterConfig::round_robin_with_retry(3303, retry_config);
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(19204, 1.0), // Fail
+                TestWorkerConfig::flaky(19205, 1.0), // Fail
+                TestWorkerConfig::healthy(19206),    // Succeed
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // With round robin and retries, should eventually hit the healthy worker
+        let payload = json!({
+            "text": "Test find healthy 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,
+            "Should succeed by retrying until finding healthy worker"
+        );
+
+        ctx.shutdown().await;
+    }
+}

third_party/sglang/sgl-model-gateway/tests/reliability_tests.rs

@@ -0,0 +1,8 @@
+//! Reliability integration tests
+
+#[path = "common/mod.rs"]
+pub mod common;
+
+mod reliability;
+
+pub use reliability::*;

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

@@ -0,0 +1,169 @@
+use std::{collections::HashMap, sync::Arc};
+
+use smg::{
+    core::{BasicWorkerBuilder, Worker, WorkerType},
+    policies::{CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, SelectWorkerInfo},
+};
+
+#[tokio::test]
+async 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()
+            },
+        )
+        .await;
+    assert!(selected.is_some(), "Should select a worker");
+
+    // Remove workers - should work without errors
+    policy.remove_worker(&worker1);
+    policy.remove_worker(&worker2);
+}
+
+#[tokio::test]
+async 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).await;
+    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).await;
+    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).await;
+    assert!(selected.is_some(), "Should select from all workers");
+}
+
+#[tokio::test]
+async 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()
+            },
+        )
+        .await;
+    assert_eq!(selected, Some(0), "Should only have worker2 left");
+}

third_party/sglang/sgl-model-gateway/tests/routing/header_forwarding_test.rs

@@ -0,0 +1,308 @@
+//! Header forwarding integration tests
+//!
+//! Tests for header propagation through the router to workers.
+
+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,
+};
+
+#[cfg(test)]
+mod header_forwarding_tests {
+    use super::*;
+
+    /// Test that X-Request-Id header is forwarded
+    #[tokio::test]
+    async fn test_request_id_forwarding() {
+        let ctx = AppTestContext::new(vec![MockWorkerConfig {
+            port: 19400,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let app = ctx.create_app().await;
+
+        let custom_request_id = "test-request-id-12345";
+        let payload = json!({
+            "text": "Test header forwarding",
+            "stream": false
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/generate")
+            .header(CONTENT_TYPE, "application/json")
+            .header("x-request-id", custom_request_id)
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        assert_eq!(resp.status(), StatusCode::OK);
+
+        // Response should have the same request ID
+        let response_id = resp.headers().get("x-request-id");
+        assert!(response_id.is_some(), "Response should have x-request-id");
+        assert_eq!(
+            response_id.unwrap().to_str().unwrap(),
+            custom_request_id,
+            "Request ID should be preserved"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test custom request ID headers
+    #[tokio::test]
+    async fn test_custom_request_id_headers() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .random_policy()
+            .host("127.0.0.1")
+            .port(3500)
+            .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)
+            .request_id_headers(vec![
+                "custom-trace-id".to_string(),
+                "x-correlation-id".to_string(),
+            ])
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![MockWorkerConfig {
+                port: 19401,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        let custom_trace_id = "my-custom-trace-123";
+        let payload = json!({
+            "text": "Test custom headers",
+            "stream": false
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/generate")
+            .header(CONTENT_TYPE, "application/json")
+            .header("custom-trace-id", custom_trace_id)
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        assert_eq!(resp.status(), StatusCode::OK);
+
+        // Response should use the custom header as request ID
+        let response_id = resp.headers().get("x-request-id");
+        assert!(response_id.is_some());
+        assert_eq!(
+            response_id.unwrap().to_str().unwrap(),
+            custom_trace_id,
+            "Custom header should be used as request ID"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test correlation ID header forwarding
+    #[tokio::test]
+    async fn test_correlation_id_forwarding() {
+        let ctx = AppTestContext::new(vec![MockWorkerConfig {
+            port: 19402,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let app = ctx.create_app().await;
+
+        let correlation_id = "correlation-abc-789";
+        let payload = json!({
+            "text": "Test correlation ID",
+            "stream": false
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/generate")
+            .header(CONTENT_TYPE, "application/json")
+            .header("x-correlation-id", correlation_id)
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        assert_eq!(resp.status(), StatusCode::OK);
+
+        // Response should preserve the correlation ID as request ID
+        let response_id = resp.headers().get("x-request-id");
+        assert!(response_id.is_some());
+        assert_eq!(
+            response_id.unwrap().to_str().unwrap(),
+            correlation_id,
+            "Correlation ID should be preserved"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test that request ID is generated when not provided
+    #[tokio::test]
+    async fn test_auto_generated_request_id() {
+        let ctx = AppTestContext::new(vec![MockWorkerConfig {
+            port: 19403,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let app = ctx.create_app().await;
+
+        let payload = json!({
+            "text": "Test auto-generated ID",
+            "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);
+
+        // Response should have an auto-generated request ID
+        let response_id = resp.headers().get("x-request-id");
+        assert!(
+            response_id.is_some(),
+            "Response should have auto-generated x-request-id"
+        );
+
+        let id_value = response_id.unwrap().to_str().unwrap();
+        assert!(!id_value.is_empty(), "Request ID should not be empty");
+        // For generate endpoint, ID should have 'gnt-' prefix
+        assert!(
+            id_value.starts_with("gnt-"),
+            "Generate endpoint should have gnt- prefix, got: {}",
+            id_value
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test chat completions request ID format
+    #[tokio::test]
+    async fn test_chat_completions_request_id_format() {
+        let ctx = AppTestContext::new(vec![MockWorkerConfig {
+            port: 19404,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let app = ctx.create_app().await;
+
+        let payload = json!({
+            "model": "test-model",
+            "messages": [{"role": "user", "content": "Hello"}],
+            "stream": false
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/v1/chat/completions")
+            .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);
+
+        // Response should have chatcmpl- prefix for chat completions
+        let response_id = resp.headers().get("x-request-id");
+        assert!(response_id.is_some());
+
+        let id_value = response_id.unwrap().to_str().unwrap();
+        assert!(
+            id_value.starts_with("chatcmpl-"),
+            "Chat completions should have chatcmpl- prefix, got: {}",
+            id_value
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test multiple header priorities
+    #[tokio::test]
+    async fn test_header_priority() {
+        let ctx = AppTestContext::new(vec![MockWorkerConfig {
+            port: 19405,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        let app = ctx.create_app().await;
+
+        let primary_id = "primary-request-id";
+        let fallback_id = "fallback-correlation-id";
+        let payload = json!({
+            "text": "Test header priority",
+            "stream": false
+        });
+
+        // When x-request-id is provided, it should take priority
+        let req = Request::builder()
+            .method("POST")
+            .uri("/generate")
+            .header(CONTENT_TYPE, "application/json")
+            .header("x-request-id", primary_id)
+            .header("x-correlation-id", fallback_id)
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        assert_eq!(resp.status(), StatusCode::OK);
+
+        let response_id = resp.headers().get("x-request-id");
+        assert!(response_id.is_some());
+        assert_eq!(
+            response_id.unwrap().to_str().unwrap(),
+            primary_id,
+            "x-request-id should take priority"
+        );
+
+        ctx.shutdown().await;
+    }
+}

third_party/sglang/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;
+    }
+}

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

@@ -0,0 +1,346 @@
+//! 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;
+    }
+}
+
+#[cfg(test)]
+mod manual_min_group_tests {
+    use super::*;
+
+    async fn send_request(app: axum::Router, routing_key: &str) -> (String, String) {
+        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.oneshot(req).await.unwrap();
+        assert_eq!(resp.status(), StatusCode::OK);
+
+        let worker_id = resp
+            .headers()
+            .get("x-worker-id")
+            .expect("Response should have x-worker-id header")
+            .to_str()
+            .unwrap()
+            .to_string();
+
+        (routing_key.to_string(), worker_id)
+    }
+
+    #[tokio::test]
+    async fn test_min_group_concurrent_distribution() {
+        let config = TestRouterConfig::manual_min_group(3910);
+
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::slow_workers(29910, 3, 500))
+                .await;
+
+        let app = ctx.create_app().await;
+
+        let mut handles = Vec::new();
+        for i in 0..9 {
+            let routing_key = format!("key-{}", i);
+            let app_clone = app.clone();
+            let handle = tokio::spawn(async move { send_request(app_clone, &routing_key).await });
+            handles.push(handle);
+        }
+
+        let results: Vec<(String, String)> = futures_util::future::join_all(handles)
+            .await
+            .into_iter()
+            .map(|r| r.unwrap())
+            .collect();
+
+        let key_to_worker: HashMap<String, String> = results.into_iter().collect();
+
+        let worker_counts: HashMap<String, usize> =
+            key_to_worker.values().fold(HashMap::new(), |mut acc, w| {
+                *acc.entry(w.clone()).or_default() += 1;
+                acc
+            });
+
+        assert_eq!(
+            worker_counts.len(),
+            3,
+            "min_group should distribute keys across all 3 workers, got {:?}",
+            worker_counts
+        );
+        for (worker, count) in &worker_counts {
+            assert_eq!(
+                *count, 3,
+                "Worker {} should have exactly 3 keys, got {}. Distribution: {:?}",
+                worker, count, key_to_worker
+            );
+        }
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_min_group_sticky_routing() {
+        let config = TestRouterConfig::manual_min_group(3911);
+
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::slow_workers(29920, 3, 200))
+                .await;
+
+        let app = ctx.create_app().await;
+
+        let routing_key = "sticky-key-123";
+
+        let mut handles = Vec::new();
+        for _ in 0..5 {
+            let app_clone = app.clone();
+            let key = routing_key.to_string();
+            let handle = tokio::spawn(async move { send_request(app_clone, &key).await });
+            handles.push(handle);
+        }
+
+        let results: Vec<(String, String)> = futures_util::future::join_all(handles)
+            .await
+            .into_iter()
+            .map(|r| r.unwrap())
+            .collect();
+
+        let workers: Vec<String> = results.into_iter().map(|(_, w)| w).collect();
+        let unique_workers: HashSet<&String> = workers.iter().collect();
+        assert_eq!(
+            unique_workers.len(),
+            1,
+            "All requests with same routing key should route to same worker, got {:?}",
+            unique_workers
+        );
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_min_group_mixed_concurrent_routing() {
+        let config = TestRouterConfig::manual_min_group(3912);
+
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::slow_workers(29930, 2, 300))
+                .await;
+
+        let app = ctx.create_app().await;
+
+        let mut handles = Vec::new();
+        for i in 0..4 {
+            let routing_key = format!("key-{}", i);
+            for _ in 0..3 {
+                let app_clone = app.clone();
+                let key = routing_key.clone();
+                let handle = tokio::spawn(async move { send_request(app_clone, &key).await });
+                handles.push(handle);
+            }
+        }
+
+        let results: Vec<(String, String)> = futures_util::future::join_all(handles)
+            .await
+            .into_iter()
+            .map(|r| r.unwrap())
+            .collect();
+
+        let mut key_to_workers: HashMap<String, HashSet<String>> = HashMap::new();
+        for (key, worker) in results {
+            key_to_workers.entry(key).or_default().insert(worker);
+        }
+
+        for (key, workers) in &key_to_workers {
+            assert_eq!(
+                workers.len(),
+                1,
+                "Key {} should route to exactly one worker (sticky), but got {:?}",
+                key,
+                workers
+            );
+        }
+
+        let all_workers: HashSet<String> = key_to_workers.values().flatten().cloned().collect();
+        assert_eq!(
+            all_workers.len(),
+            2,
+            "Keys should be distributed across both workers"
+        );
+
+        ctx.shutdown().await;
+    }
+}

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

@@ -0,0 +1,14 @@
+//! Routing integration tests
+
+pub mod cache_aware_backward_compat_test;
+pub mod header_forwarding_test;
+pub mod load_balancing_test;
+pub mod manual_routing_test;
+pub mod payload_size_test;
+pub mod pd_routing_test;
+pub mod policy_registry_integration;
+pub mod power_of_two_test;
+pub mod service_discovery_test;
+pub mod test_openai_routing;
+pub mod test_pd_routing;
+pub mod worker_management_test;

third_party/sglang/sgl-model-gateway/tests/routing/payload_size_test.rs

@@ -0,0 +1,297 @@
+//! Payload size integration tests
+//!
+//! Tests for request payload size limits and handling.
+
+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,
+};
+
+#[cfg(test)]
+mod payload_size_tests {
+    use super::*;
+
+    /// Test that small payloads are handled correctly
+    #[tokio::test]
+    async fn test_small_payload() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4200)
+            .max_payload_size(1024 * 1024) // 1MB limit
+            .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![MockWorkerConfig {
+                port: 20200,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        let payload = json!({
+            "text": "Small payload test",
+            "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,
+            "Small payload should be accepted"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test that payloads within limit are accepted
+    #[tokio::test]
+    async fn test_payload_within_limit() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4201)
+            .max_payload_size(1024 * 1024) // 1MB limit
+            .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![MockWorkerConfig {
+                port: 20201,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Create a ~100KB payload (well within 1MB limit)
+        let large_text = "x".repeat(100 * 1024);
+        let payload = json!({
+            "text": large_text,
+            "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,
+            "Payload within limit should be accepted"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test that payloads exceeding limit are rejected
+    #[tokio::test]
+    async fn test_payload_exceeds_limit() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4202)
+            .max_payload_size(1024) // Very small 1KB limit
+            .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![MockWorkerConfig {
+                port: 20202,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Create a payload larger than 1KB limit
+        let large_text = "x".repeat(2048);
+        let payload = json!({
+            "text": large_text,
+            "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 be rejected with 413 Payload Too Large or similar
+        assert!(
+            resp.status() == StatusCode::PAYLOAD_TOO_LARGE
+                || resp.status() == StatusCode::BAD_REQUEST,
+            "Payload exceeding limit should be rejected, got {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test edge case: payload exactly at limit
+    #[tokio::test]
+    async fn test_payload_at_exact_limit() {
+        // Use a more reasonable limit for this test
+        let limit_bytes = 10 * 1024; // 10KB limit
+
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4203)
+            .max_payload_size(limit_bytes)
+            .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![MockWorkerConfig {
+                port: 20203,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Create a payload slightly under the limit (accounting for JSON overhead)
+        let text_size = limit_bytes - 100; // Leave room for JSON structure
+        let text = "x".repeat(text_size);
+        let payload = json!({
+            "text": text,
+            "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();
+        // Payload at/near limit should be accepted
+        assert!(
+            resp.status() == StatusCode::OK || resp.status() == StatusCode::PAYLOAD_TOO_LARGE,
+            "Payload at limit boundary, got status {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test default payload size limit (256MB)
+    #[tokio::test]
+    async fn test_default_payload_limit() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4204)
+            .max_payload_size(256 * 1024 * 1024) // Default 256MB
+            .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![MockWorkerConfig {
+                port: 20204,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Create a 1MB payload (well within 256MB)
+        let large_text = "x".repeat(1024 * 1024);
+        let payload = json!({
+            "text": large_text,
+            "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,
+            "1MB payload should be accepted with 256MB limit"
+        );
+
+        ctx.shutdown().await;
+    }
+}

third_party/sglang/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;
+    }
+}

third_party/sglang/sgl-model-gateway/tests/routing/policy_registry_integration.rs

@@ -0,0 +1,173 @@
+//! Integration tests for PolicyRegistry with RouterManager
+
+use std::{collections::HashMap, sync::Arc};
+
+use smg::{
+    config::PolicyConfig, core::WorkerRegistry, policies::PolicyRegistry,
+    protocols::worker_spec::WorkerConfigRequest, routers::router_manager::RouterManager,
+};
+
+#[tokio::test]
+async fn test_policy_registry_with_router_manager() {
+    // Create HTTP client
+    let _client = reqwest::Client::new();
+
+    // Create shared registries
+    let worker_registry = Arc::new(WorkerRegistry::new());
+    let policy_registry = Arc::new(PolicyRegistry::new(PolicyConfig::RoundRobin));
+
+    // Create RouterManager with shared registries
+    let _router_manager = RouterManager::new(worker_registry.clone());
+
+    // Add first worker for llama-3 with cache_aware policy hint
+    let mut labels1 = HashMap::new();
+    labels1.insert("policy".to_string(), "cache_aware".to_string());
+
+    let _worker1_config = WorkerConfigRequest {
+        url: "http://worker1:8000".to_string(),
+        model_id: Some("llama-3".to_string()),
+        api_key: Some("test_api_key".to_string()),
+        worker_type: None,
+        priority: None,
+        cost: None,
+        labels: labels1,
+        bootstrap_port: None,
+        tokenizer_path: None,
+        reasoning_parser: None,
+        tool_parser: None,
+        chat_template: None,
+        runtime: None,
+        health_check_timeout_secs: 30,
+        health_check_interval_secs: 60,
+        health_success_threshold: 2,
+        health_failure_threshold: 3,
+        disable_health_check: false,
+        max_connection_attempts: 20,
+        dp_aware: false,
+    };
+
+    // This would normally connect to a real worker, but for testing we'll just verify the structure
+    // In a real test, we'd need to mock the worker or use a test server
+
+    let _llama_policy = policy_registry.get_policy("llama-3");
+    // After first worker is added, llama-3 should have a policy
+
+    // Add second worker for llama-3 with different policy hint (should be ignored)
+    let mut labels2 = HashMap::new();
+    labels2.insert("policy".to_string(), "random".to_string());
+
+    let _worker2_config = WorkerConfigRequest {
+        url: "http://worker2:8000".to_string(),
+        model_id: Some("llama-3".to_string()),
+        api_key: Some("test_api_key".to_string()),
+        worker_type: None,
+        priority: None,
+        cost: None,
+        labels: labels2,
+        bootstrap_port: None,
+        tokenizer_path: None,
+        reasoning_parser: None,
+        tool_parser: None,
+        chat_template: None,
+        runtime: None,
+        health_check_timeout_secs: 30,
+        health_check_interval_secs: 60,
+        health_success_threshold: 2,
+        health_failure_threshold: 3,
+        disable_health_check: false,
+        max_connection_attempts: 20,
+        dp_aware: false,
+    };
+
+    // The second worker should use the same policy as the first (cache_aware)
+
+    // Add worker for different model (gpt-4) with random policy
+    let mut labels3 = HashMap::new();
+    labels3.insert("policy".to_string(), "random".to_string());
+
+    let _worker3_config = WorkerConfigRequest {
+        url: "http://worker3:8000".to_string(),
+        model_id: Some("gpt-4".to_string()),
+        api_key: Some("test_api_key".to_string()),
+        worker_type: None,
+        priority: None,
+        cost: None,
+        labels: labels3,
+        bootstrap_port: None,
+        tokenizer_path: None,
+        reasoning_parser: None,
+        tool_parser: None,
+        runtime: None,
+        chat_template: None,
+        health_check_timeout_secs: 30,
+        health_check_interval_secs: 60,
+        health_success_threshold: 2,
+        health_failure_threshold: 3,
+        disable_health_check: false,
+        max_connection_attempts: 20,
+        dp_aware: false,
+    };
+
+    let _gpt_policy = policy_registry.get_policy("gpt-4");
+
+    // When we remove both llama-3 workers, the policy should be cleaned up
+
+    println!("PolicyRegistry integration test structure created");
+    println!("Note: This test requires mocking or test servers to fully execute");
+}
+
+#[test]
+fn test_policy_registry_cleanup() {
+    use smg::{config::PolicyConfig, policies::PolicyRegistry};
+
+    let registry = PolicyRegistry::new(PolicyConfig::RoundRobin);
+
+    // Add workers for a model
+    let policy1 = registry.on_worker_added("model-1", Some("cache_aware"));
+    assert_eq!(policy1.name(), "cache_aware");
+
+    // Second worker uses existing policy
+    let policy2 = registry.on_worker_added("model-1", Some("random"));
+    assert_eq!(policy2.name(), "cache_aware"); // Should still be cache_aware
+
+    assert!(registry.get_policy("model-1").is_some());
+
+    // Remove first worker - policy should remain
+    registry.on_worker_removed("model-1");
+    assert!(registry.get_policy("model-1").is_some());
+
+    // Remove second worker - policy should be cleaned up
+    registry.on_worker_removed("model-1");
+    assert!(registry.get_policy("model-1").is_none());
+
+    println!("✓ PolicyRegistry cleanup test passed");
+}
+
+#[test]
+fn test_policy_registry_multiple_models() {
+    use smg::{config::PolicyConfig, policies::PolicyRegistry};
+
+    let registry = PolicyRegistry::new(PolicyConfig::RoundRobin);
+
+    // Add workers for different models with different policies
+    let llama_policy = registry.on_worker_added("llama-3", Some("cache_aware"));
+    let gpt_policy = registry.on_worker_added("gpt-4", Some("random"));
+    let mistral_policy = registry.on_worker_added("mistral", None); // Uses default
+
+    assert_eq!(llama_policy.name(), "cache_aware");
+    assert_eq!(gpt_policy.name(), "random");
+    assert_eq!(mistral_policy.name(), "round_robin"); // Default
+
+    assert!(registry.get_policy("llama-3").is_some());
+    assert!(registry.get_policy("gpt-4").is_some());
+    assert!(registry.get_policy("mistral").is_some());
+
+    // Get all mappings
+    let mappings = registry.get_all_mappings();
+    assert_eq!(mappings.len(), 3);
+    assert_eq!(mappings.get("llama-3").unwrap(), "cache_aware");
+    assert_eq!(mappings.get("gpt-4").unwrap(), "random");
+    assert_eq!(mappings.get("mistral").unwrap(), "round_robin");
+
+    println!("✓ PolicyRegistry multiple models test passed");
+}

third_party/sglang/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;
+    }
+}

third_party/sglang/sgl-model-gateway/tests/routing/service_discovery_test.rs

@@ -0,0 +1,260 @@
+//! Service discovery integration tests
+//!
+//! Tests for service discovery shim functionality for dynamic worker registration.
+
+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,
+};
+
+#[cfg(test)]
+mod service_discovery_tests {
+    use super::*;
+
+    /// Test service discovery endpoint responds correctly
+    #[tokio::test]
+    async fn test_service_discovery_endpoint() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4000)
+            .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![MockWorkerConfig {
+                port: 20000,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Check if service discovery endpoint exists
+        let req = Request::builder()
+            .method("GET")
+            .uri("/v1/workers")
+            .body(Body::empty())
+            .unwrap();
+
+        let resp = app.clone().oneshot(req).await.unwrap();
+        // Endpoint might return OK with worker list or 404 if not implemented
+        assert!(
+            resp.status() == StatusCode::OK || resp.status() == StatusCode::NOT_FOUND,
+            "Workers endpoint should respond, got {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test worker registration via discovery shim
+    #[tokio::test]
+    async fn test_worker_registration() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4001)
+            .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![MockWorkerConfig {
+                port: 20001,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Register a new worker via discovery endpoint
+        let register_payload = json!({
+            "url": "http://127.0.0.1:20002",
+            "weight": 1.0
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/register_worker")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(
+                serde_json::to_string(&register_payload).unwrap(),
+            ))
+            .unwrap();
+
+        let resp = app.clone().oneshot(req).await.unwrap();
+        // Registration might succeed or endpoint might not exist
+        assert!(
+            resp.status() == StatusCode::OK
+                || resp.status() == StatusCode::ACCEPTED
+                || resp.status() == StatusCode::NOT_FOUND,
+            "Registration should respond appropriately, got {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test worker deregistration via discovery shim
+    #[tokio::test]
+    async fn test_worker_deregistration() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4002)
+            .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![
+                MockWorkerConfig {
+                    port: 20003,
+                    worker_type: WorkerType::Regular,
+                    health_status: HealthStatus::Healthy,
+                    response_delay_ms: 0,
+                    fail_rate: 0.0,
+                },
+                MockWorkerConfig {
+                    port: 20004,
+                    worker_type: WorkerType::Regular,
+                    health_status: HealthStatus::Healthy,
+                    response_delay_ms: 0,
+                    fail_rate: 0.0,
+                },
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Deregister a worker via discovery endpoint
+        let deregister_payload = json!({
+            "url": "http://127.0.0.1:20003"
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/deregister_worker")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(
+                serde_json::to_string(&deregister_payload).unwrap(),
+            ))
+            .unwrap();
+
+        let resp = app.clone().oneshot(req).await.unwrap();
+        assert!(
+            resp.status() == StatusCode::OK
+                || resp.status() == StatusCode::ACCEPTED
+                || resp.status() == StatusCode::NOT_FOUND,
+            "Deregistration should respond appropriately, got {}",
+            resp.status()
+        );
+
+        // Requests should still work with remaining worker
+        let payload = json!({
+            "text": "Test after deregistration",
+            "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 with remaining worker"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test health status reporting for discovery
+    #[tokio::test]
+    async fn test_health_status_endpoint() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4003)
+            .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![MockWorkerConfig {
+                port: 20005,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Check health endpoint
+        let req = Request::builder()
+            .method("GET")
+            .uri("/health")
+            .body(Body::empty())
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Health endpoint should return OK when workers are healthy"
+        );
+
+        ctx.shutdown().await;
+    }
+}

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

@@ -0,0 +1,906 @@
+//! 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 data_connector::{ResponseId, StoredResponse};
+use serde_json::json;
+use smg::{
+    config::{ConfigError, HistoryBackend, OracleConfig, RouterConfig, RoutingMode},
+    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 = config
+        .validate()
+        .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();
+
+    config.validate().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();
+
+    config
+        .validate()
+        .expect("wallet-based config should validate");
+}

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

@@ -0,0 +1,926 @@
+#[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 data_connector::{
+                    MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage,
+                };
+                use smg::{
+                    core::{LoadMonitor, WorkerRegistry},
+                    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 engines, and mcp manager
+                let worker_job_queue = Arc::new(OnceLock::new());
+                let workflow_engines = 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_engines(workflow_engines)
+                        .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
+            ("/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,
+        };
+    }
+}

third_party/sglang/sgl-model-gateway/tests/routing/worker_management_test.rs

@@ -0,0 +1,139 @@
+//! Worker management integration tests
+//!
+//! Tests for dynamic worker add/remove operations via management API.
+//! The actual worker management API uses:
+//! - POST /workers - create a worker
+//! - GET /workers - list workers
+//! - DELETE /workers/{worker_id} - remove a worker
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+#[cfg(test)]
+mod worker_management_tests {
+    use super::*;
+
+    /// Test listing workers via API
+    #[tokio::test]
+    async fn test_list_workers() {
+        let config = TestRouterConfig::round_robin(3900);
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::healthy(19900),
+                TestWorkerConfig::healthy(19901),
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // List workers via GET /workers
+        let req = Request::builder()
+            .method("GET")
+            .uri("/workers")
+            .body(Body::empty())
+            .unwrap();
+
+        let resp = app.clone().oneshot(req).await.unwrap();
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "GET /workers should return OK"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test that routing continues to work with multiple workers
+    #[tokio::test]
+    async fn test_routing_with_multiple_workers() {
+        let config = TestRouterConfig::round_robin(3901);
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::healthy(19902),
+                TestWorkerConfig::healthy(19903),
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+        let mut success_count = 0;
+
+        // Verify routing distributes across workers
+        for i in 0..20 {
+            let payload = json!({
+                "text": format!("Test 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;
+            }
+        }
+
+        assert_eq!(
+            success_count, 20,
+            "All requests should succeed with multiple workers"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test that requests continue to work during worker operations
+    #[tokio::test]
+    async fn test_requests_during_worker_changes() {
+        let config = TestRouterConfig::round_robin(3902);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(19904)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Send requests and verify they succeed
+        let mut success_count = 0;
+        for i in 0..10 {
+            let payload = json!({
+                "text": format!("Request during changes {}", 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, 10,
+            "All requests should succeed during normal operation"
+        );
+
+        ctx.shutdown().await;
+    }
+}

third_party/sglang/sgl-model-gateway/tests/routing_tests.rs

@@ -0,0 +1,8 @@
+//! Routing integration tests
+
+#[path = "common/mod.rs"]
+pub mod common;
+
+mod routing;
+
+pub use routing::*;

third_party/sglang/sgl-model-gateway/tests/security/auth_integration_test.rs

@@ -0,0 +1,937 @@
+//! Integration tests for control plane authentication.
+//!
+//! Tests the full authentication flow including:
+//! - JWT token validation with mock JWKS server
+//! - API key authentication
+//! - Role-based access control
+//! - Token expiration and replay protection
+
+use std::{
+    net::SocketAddr,
+    sync::LazyLock,
+    time::{Duration, SystemTime, UNIX_EPOCH},
+};
+
+use axum::{routing::get, Json, Router};
+use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
+use jsonwebtoken::{encode, EncodingKey, Header};
+use rsa::{traits::PublicKeyParts, RsaPrivateKey};
+use serde::{Deserialize, Serialize};
+use serde_json::json;
+use smg::auth::{ApiKeyEntry, ControlPlaneAuthConfig, ControlPlaneAuthState, JwtConfig, Role};
+use tokio::net::TcpListener;
+
+const TEST_KEY_ID: &str = "test-key-1";
+
+/// Test RSA key pair - generates matching public/private key components
+struct TestKeyPair {
+    private_key_pem: String,
+    n_base64url: String,
+    e_base64url: String,
+}
+
+impl TestKeyPair {
+    fn generate() -> Self {
+        // Generate a new 2048-bit RSA key pair for testing
+        use rsa::{pkcs8::EncodePrivateKey, rand_core::OsRng};
+
+        let mut rng = OsRng;
+        let private_key =
+            RsaPrivateKey::new(&mut rng, 2048).expect("Failed to generate RSA key pair");
+
+        // Get the public key components
+        let n_bytes = private_key.n().to_bytes_be();
+        let e_bytes = private_key.e().to_bytes_be();
+
+        // Encode to base64url (no padding) for JWKS
+        let n_base64url = URL_SAFE_NO_PAD.encode(&n_bytes);
+        let e_base64url = URL_SAFE_NO_PAD.encode(&e_bytes);
+
+        // Export private key as PEM
+        let private_key_pem = private_key
+            .to_pkcs8_pem(rsa::pkcs8::LineEnding::LF)
+            .expect("Failed to export private key")
+            .to_string();
+
+        Self {
+            private_key_pem,
+            n_base64url,
+            e_base64url,
+        }
+    }
+}
+
+/// Lazily initialized test key pair (shared across all tests)
+static TEST_KEYS: LazyLock<TestKeyPair> = LazyLock::new(TestKeyPair::generate);
+
+/// Claims structure for test tokens
+#[derive(Debug, Serialize, Deserialize)]
+struct TestClaims {
+    sub: String,
+    iss: String,
+    aud: String,
+    exp: u64,
+    iat: u64,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    jti: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    name: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    roles: Option<Vec<String>>,
+}
+
+/// Create a test JWT token using the generated test key pair
+fn create_test_token(claims: &TestClaims) -> String {
+    let mut header = Header::new(jsonwebtoken::Algorithm::RS256);
+    header.kid = Some(TEST_KEY_ID.to_string());
+
+    let key = EncodingKey::from_rsa_pem(TEST_KEYS.private_key_pem.as_bytes())
+        .expect("Failed to create encoding key");
+
+    encode(&header, claims, &key).expect("Failed to encode token")
+}
+
+/// Create claims with default values
+fn create_claims(sub: &str, roles: Vec<&str>) -> TestClaims {
+    let now = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_secs();
+
+    TestClaims {
+        sub: sub.to_string(),
+        iss: "https://test-issuer.example.com".to_string(),
+        aud: "test-gateway".to_string(),
+        exp: now + 3600, // 1 hour from now
+        iat: now,
+        jti: Some(uuid::Uuid::new_v4().to_string()),
+        name: Some("Test User".to_string()),
+        roles: Some(roles.into_iter().map(String::from).collect()),
+    }
+}
+
+/// Start a mock JWKS server using the generated test key pair
+async fn start_mock_jwks_server() -> (SocketAddr, tokio::task::JoinHandle<()>) {
+    // Use the generated key pair components
+    let jwks_response = json!({
+        "keys": [{
+            "kty": "RSA",
+            "use": "sig",
+            "kid": TEST_KEY_ID,
+            "alg": "RS256",
+            "n": TEST_KEYS.n_base64url,
+            "e": TEST_KEYS.e_base64url
+        }]
+    });
+
+    let app = Router::new().route(
+        "/.well-known/jwks.json",
+        get(move || {
+            let jwks = jwks_response.clone();
+            async move { Json(jwks) }
+        }),
+    );
+
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+
+    let handle = tokio::spawn(async move {
+        axum::serve(listener, app).await.unwrap();
+    });
+
+    // Give server time to start
+    tokio::time::sleep(Duration::from_millis(50)).await;
+
+    (addr, handle)
+}
+
+// ============================================================================
+// JWT Authentication Tests
+// ============================================================================
+
+#[tokio::test]
+async fn test_jwt_valid_token_with_admin_role() {
+    let (addr, _server) = start_mock_jwks_server().await;
+
+    let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway")
+        .with_jwks_uri(format!(
+            "http://127.0.0.1:{}/.well-known/jwks.json",
+            addr.port()
+        ))
+        .with_role_mapping("admin", Role::Admin);
+
+    let config = ControlPlaneAuthConfig {
+        jwt: Some(jwt_config),
+        api_keys: vec![],
+        audit_enabled: false,
+    };
+
+    let state = ControlPlaneAuthState::from_config(config)
+        .await
+        .expect("Failed to create auth state");
+
+    // Create a valid token with admin role
+    let claims = create_claims("admin@example.com", vec!["admin"]);
+    let token = create_test_token(&claims);
+
+    // Validate the token
+    let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set");
+    let result = jwt_validator.validate(&token).await;
+
+    assert!(
+        result.is_ok(),
+        "Token validation failed: {:?}",
+        result.err()
+    );
+    let validated = result.unwrap();
+    assert_eq!(validated.subject, "admin@example.com");
+    assert_eq!(validated.role, Role::Admin);
+}
+
+#[tokio::test]
+async fn test_jwt_valid_token_with_user_role() {
+    let (addr, _server) = start_mock_jwks_server().await;
+
+    let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway")
+        .with_jwks_uri(format!(
+            "http://127.0.0.1:{}/.well-known/jwks.json",
+            addr.port()
+        ))
+        .with_role_mapping("admin", Role::Admin);
+
+    let config = ControlPlaneAuthConfig {
+        jwt: Some(jwt_config),
+        api_keys: vec![],
+        audit_enabled: false,
+    };
+
+    let state = ControlPlaneAuthState::from_config(config)
+        .await
+        .expect("Failed to create auth state");
+
+    // Create a valid token with user role (not admin)
+    let claims = create_claims("user@example.com", vec!["user", "viewer"]);
+    let token = create_test_token(&claims);
+
+    let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set");
+    let result = jwt_validator.validate(&token).await;
+
+    assert!(result.is_ok());
+    let validated = result.unwrap();
+    assert_eq!(validated.subject, "user@example.com");
+    assert_eq!(validated.role, Role::User); // Not admin
+}
+
+#[tokio::test]
+async fn test_jwt_expired_token() {
+    let (addr, _server) = start_mock_jwks_server().await;
+
+    let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway")
+        .with_jwks_uri(format!(
+            "http://127.0.0.1:{}/.well-known/jwks.json",
+            addr.port()
+        ));
+
+    let config = ControlPlaneAuthConfig {
+        jwt: Some(jwt_config),
+        api_keys: vec![],
+        audit_enabled: false,
+    };
+
+    let state = ControlPlaneAuthState::from_config(config)
+        .await
+        .expect("Failed to create auth state");
+
+    // Create an expired token
+    let now = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_secs();
+
+    let claims = TestClaims {
+        sub: "user@example.com".to_string(),
+        iss: "https://test-issuer.example.com".to_string(),
+        aud: "test-gateway".to_string(),
+        exp: now - 3600, // Expired 1 hour ago
+        iat: now - 7200,
+        jti: None,
+        name: None,
+        roles: Some(vec!["admin".to_string()]),
+    };
+    let token = create_test_token(&claims);
+
+    let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set");
+    let result = jwt_validator.validate(&token).await;
+
+    assert!(result.is_err(), "Expired token should fail validation");
+    let err = result.unwrap_err();
+    assert!(
+        err.to_string().contains("exp") || err.to_string().contains("ExpiredSignature"),
+        "Error should mention expiration: {}",
+        err
+    );
+}
+
+#[tokio::test]
+async fn test_jwt_wrong_audience() {
+    let (addr, _server) = start_mock_jwks_server().await;
+
+    let jwt_config =
+        JwtConfig::new("https://test-issuer.example.com", "correct-audience").with_jwks_uri(
+            format!("http://127.0.0.1:{}/.well-known/jwks.json", addr.port()),
+        );
+
+    let config = ControlPlaneAuthConfig {
+        jwt: Some(jwt_config),
+        api_keys: vec![],
+        audit_enabled: false,
+    };
+
+    let state = ControlPlaneAuthState::from_config(config)
+        .await
+        .expect("Failed to create auth state");
+
+    // Create a token with wrong audience
+    let now = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_secs();
+
+    let claims = TestClaims {
+        sub: "user@example.com".to_string(),
+        iss: "https://test-issuer.example.com".to_string(),
+        aud: "wrong-audience".to_string(), // Wrong audience
+        exp: now + 3600,
+        iat: now,
+        jti: None,
+        name: None,
+        roles: Some(vec!["admin".to_string()]),
+    };
+    let token = create_test_token(&claims);
+
+    let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set");
+    let result = jwt_validator.validate(&token).await;
+
+    assert!(result.is_err(), "Wrong audience should fail validation");
+}
+
+#[tokio::test]
+async fn test_jwt_wrong_issuer() {
+    let (addr, _server) = start_mock_jwks_server().await;
+
+    let jwt_config =
+        JwtConfig::new("https://correct-issuer.example.com", "test-gateway").with_jwks_uri(
+            format!("http://127.0.0.1:{}/.well-known/jwks.json", addr.port()),
+        );
+
+    let config = ControlPlaneAuthConfig {
+        jwt: Some(jwt_config),
+        api_keys: vec![],
+        audit_enabled: false,
+    };
+
+    let state = ControlPlaneAuthState::from_config(config)
+        .await
+        .expect("Failed to create auth state");
+
+    // Create a token with wrong issuer
+    let claims = create_claims("user@example.com", vec!["admin"]);
+    // claims has issuer "https://test-issuer.example.com" which doesn't match
+    let token = create_test_token(&claims);
+
+    let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set");
+    let result = jwt_validator.validate(&token).await;
+
+    assert!(result.is_err(), "Wrong issuer should fail validation");
+}
+
+// ============================================================================
+// API Key Authentication Tests
+// ============================================================================
+
+#[tokio::test]
+async fn test_api_key_valid_admin() {
+    let config = ControlPlaneAuthConfig {
+        jwt: None,
+        api_keys: vec![ApiKeyEntry::new(
+            "admin-key-1",
+            "Admin Service Account",
+            "sk-test-admin-key-12345",
+            Role::Admin,
+        )],
+        audit_enabled: false,
+    };
+
+    // Verify we can find the API key
+    let found = config.find_api_key("sk-test-admin-key-12345");
+    assert!(found.is_some(), "API key should be found");
+
+    let entry = found.unwrap();
+    assert_eq!(entry.id, "admin-key-1");
+    assert_eq!(entry.role, Role::Admin);
+}
+
+#[tokio::test]
+async fn test_api_key_valid_user() {
+    let config = ControlPlaneAuthConfig {
+        jwt: None,
+        api_keys: vec![ApiKeyEntry::new(
+            "user-key-1",
+            "User Service Account",
+            "sk-test-user-key-12345",
+            Role::User,
+        )],
+        audit_enabled: false,
+    };
+
+    let found = config.find_api_key("sk-test-user-key-12345");
+    assert!(found.is_some());
+
+    let entry = found.unwrap();
+    assert_eq!(entry.role, Role::User);
+}
+
+#[tokio::test]
+async fn test_api_key_invalid() {
+    let config = ControlPlaneAuthConfig {
+        jwt: None,
+        api_keys: vec![ApiKeyEntry::new(
+            "admin-key-1",
+            "Admin Service Account",
+            "sk-correct-key",
+            Role::Admin,
+        )],
+        audit_enabled: false,
+    };
+
+    // Try with wrong key
+    let found = config.find_api_key("sk-wrong-key");
+    assert!(found.is_none(), "Wrong API key should not be found");
+}
+
+#[tokio::test]
+async fn test_api_key_timing_attack_resistance() {
+    // Verify that key comparison uses constant-time comparison
+    let config = ControlPlaneAuthConfig {
+        jwt: None,
+        api_keys: vec![ApiKeyEntry::new(
+            "key-1",
+            "Test Key",
+            "sk-abcdefghijklmnop",
+            Role::Admin,
+        )],
+        audit_enabled: false,
+    };
+
+    // These should all take roughly the same time due to constant-time comparison
+    // (We can't easily test timing in unit tests, but we verify the function works)
+    assert!(config.find_api_key("sk-abcdefghijklmnop").is_some());
+    assert!(config.find_api_key("sk-abcdefghijklmnox").is_none()); // One char different
+    assert!(config.find_api_key("sk-xxxxxxxxxxxxxxxx").is_none()); // All different
+    assert!(config.find_api_key("sk-a").is_none()); // Much shorter
+}
+
+// ============================================================================
+// Combined Auth Tests (JWT + API Key fallback)
+// ============================================================================
+
+#[tokio::test]
+async fn test_combined_auth_jwt_and_api_keys() {
+    let (addr, _server) = start_mock_jwks_server().await;
+
+    let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway")
+        .with_jwks_uri(format!(
+            "http://127.0.0.1:{}/.well-known/jwks.json",
+            addr.port()
+        ))
+        .with_role_mapping("admin", Role::Admin);
+
+    let config = ControlPlaneAuthConfig {
+        jwt: Some(jwt_config),
+        api_keys: vec![ApiKeyEntry::new(
+            "backup-key",
+            "Backup Admin Key",
+            "sk-backup-admin-key",
+            Role::Admin,
+        )],
+        audit_enabled: true,
+    };
+
+    let state = ControlPlaneAuthState::from_config(config.clone())
+        .await
+        .expect("Failed to create auth state");
+
+    // Both JWT and API key should work
+    assert!(state.jwt_validator.is_some());
+    assert!(config.find_api_key("sk-backup-admin-key").is_some());
+
+    // Verify JWT works
+    let claims = create_claims("jwt-user@example.com", vec!["admin"]);
+    let token = create_test_token(&claims);
+    let result = state.jwt_validator.as_ref().unwrap().validate(&token).await;
+    assert!(result.is_ok());
+}
+
+// ============================================================================
+// Role-Based Access Control Tests
+// ============================================================================
+
+#[tokio::test]
+async fn test_role_admin_has_access() {
+    assert!(Role::Admin.is_admin());
+}
+
+#[tokio::test]
+async fn test_role_user_no_admin_access() {
+    assert!(!Role::User.is_admin());
+}
+
+#[tokio::test]
+async fn test_role_parsing() {
+    assert_eq!("admin".parse::<Role>().unwrap(), Role::Admin);
+    assert_eq!("Admin".parse::<Role>().unwrap(), Role::Admin);
+    assert_eq!("ADMIN".parse::<Role>().unwrap(), Role::Admin);
+    assert_eq!("user".parse::<Role>().unwrap(), Role::User);
+    assert_eq!("User".parse::<Role>().unwrap(), Role::User);
+    assert_eq!("USER".parse::<Role>().unwrap(), Role::User);
+}
+
+// ============================================================================
+// try_init Helper Tests
+// ============================================================================
+
+#[tokio::test]
+async fn test_try_init_with_no_config() {
+    let result = ControlPlaneAuthState::try_init(None).await;
+    assert!(result.is_none());
+}
+
+#[tokio::test]
+async fn test_try_init_with_disabled_config() {
+    let config = ControlPlaneAuthConfig {
+        jwt: None,
+        api_keys: vec![], // No auth configured = disabled
+        audit_enabled: false,
+    };
+
+    let result = ControlPlaneAuthState::try_init(Some(&config)).await;
+    assert!(result.is_none());
+}
+
+#[tokio::test]
+async fn test_try_init_with_api_keys_only() {
+    let config = ControlPlaneAuthConfig {
+        jwt: None,
+        api_keys: vec![ApiKeyEntry::new("key-1", "Test", "sk-test", Role::Admin)],
+        audit_enabled: false,
+    };
+
+    let result = ControlPlaneAuthState::try_init(Some(&config)).await;
+    assert!(result.is_some());
+}
+
+// ============================================================================
+// Audit Logging Tests
+// ============================================================================
+
+#[tokio::test]
+async fn test_audit_logging_enabled() {
+    let config = ControlPlaneAuthConfig {
+        jwt: None,
+        api_keys: vec![ApiKeyEntry::new("key-1", "Test", "sk-test", Role::Admin)],
+        audit_enabled: true,
+    };
+
+    let state = ControlPlaneAuthState::from_config(config)
+        .await
+        .expect("Failed to create auth state");
+
+    assert!(state.audit_logger.is_enabled());
+}
+
+#[tokio::test]
+async fn test_audit_logging_disabled() {
+    let config = ControlPlaneAuthConfig {
+        jwt: None,
+        api_keys: vec![ApiKeyEntry::new("key-1", "Test", "sk-test", Role::Admin)],
+        audit_enabled: false,
+    };
+
+    let state = ControlPlaneAuthState::from_config(config)
+        .await
+        .expect("Failed to create auth state");
+
+    assert!(!state.audit_logger.is_enabled());
+}
+
+// ============================================================================
+// JTI Replay Protection Tests
+// ============================================================================
+
+#[tokio::test]
+async fn test_jwt_jti_replay_protection() {
+    use smg::auth::JwtValidator;
+
+    let (addr, _server) = start_mock_jwks_server().await;
+
+    let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway")
+        .with_jwks_uri(format!(
+            "http://127.0.0.1:{}/.well-known/jwks.json",
+            addr.port()
+        ))
+        .with_role_mapping("admin", Role::Admin);
+
+    // Create validator with JTI replay protection enabled
+    let validator = JwtValidator::from_config_with_options(jwt_config, true)
+        .await
+        .expect("Failed to create JWT validator");
+
+    // Create a token with a specific JTI
+    let claims = create_claims("user@example.com", vec!["admin"]);
+    let token = create_test_token(&claims);
+
+    // First validation should succeed
+    let result1 = validator.validate(&token).await;
+    assert!(result1.is_ok(), "First validation should succeed");
+
+    // Second validation with same token should fail (replay)
+    let result2 = validator.validate(&token).await;
+    assert!(
+        result2.is_err(),
+        "Second validation should fail (replay detected)"
+    );
+    let err = result2.unwrap_err();
+    assert!(
+        err.to_string().contains("replay") || err.to_string().contains("already been used"),
+        "Error should mention replay: {}",
+        err
+    );
+}
+
+#[tokio::test]
+async fn test_jwt_different_tokens_no_replay() {
+    use smg::auth::JwtValidator;
+
+    let (addr, _server) = start_mock_jwks_server().await;
+
+    let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway")
+        .with_jwks_uri(format!(
+            "http://127.0.0.1:{}/.well-known/jwks.json",
+            addr.port()
+        ))
+        .with_role_mapping("admin", Role::Admin);
+
+    // Create validator with JTI replay protection enabled
+    let validator = JwtValidator::from_config_with_options(jwt_config, true)
+        .await
+        .expect("Failed to create JWT validator");
+
+    // Create two different tokens (different JTIs)
+    let claims1 = create_claims("user1@example.com", vec!["admin"]);
+    let claims2 = create_claims("user2@example.com", vec!["admin"]);
+    let token1 = create_test_token(&claims1);
+    let token2 = create_test_token(&claims2);
+
+    // Both should succeed since they have different JTIs
+    let result1 = validator.validate(&token1).await;
+    assert!(result1.is_ok(), "First token validation should succeed");
+
+    let result2 = validator.validate(&token2).await;
+    assert!(
+        result2.is_ok(),
+        "Second token validation should succeed (different JTI)"
+    );
+}
+
+// ============================================================================
+// Malformed Token Tests
+// ============================================================================
+
+#[tokio::test]
+async fn test_jwt_malformed_token() {
+    let (addr, _server) = start_mock_jwks_server().await;
+
+    let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway")
+        .with_jwks_uri(format!(
+            "http://127.0.0.1:{}/.well-known/jwks.json",
+            addr.port()
+        ));
+
+    let config = ControlPlaneAuthConfig {
+        jwt: Some(jwt_config),
+        api_keys: vec![],
+        audit_enabled: false,
+    };
+
+    let state = ControlPlaneAuthState::from_config(config)
+        .await
+        .expect("Failed to create auth state");
+
+    let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set");
+
+    // Test various malformed tokens
+    let malformed_tokens = [
+        "not-a-jwt",
+        "only.two.parts.is.not.enough",
+        "",
+        "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9",  // Only header
+        "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.", // Header with dot
+        ".....",
+    ];
+
+    for token in malformed_tokens {
+        let result = jwt_validator.validate(token).await;
+        assert!(
+            result.is_err(),
+            "Malformed token '{}' should fail validation",
+            token
+        );
+    }
+}
+
+#[tokio::test]
+async fn test_jwt_missing_kid_in_header() {
+    let (addr, _server) = start_mock_jwks_server().await;
+
+    let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway")
+        .with_jwks_uri(format!(
+            "http://127.0.0.1:{}/.well-known/jwks.json",
+            addr.port()
+        ));
+
+    let config = ControlPlaneAuthConfig {
+        jwt: Some(jwt_config),
+        api_keys: vec![],
+        audit_enabled: false,
+    };
+
+    let state = ControlPlaneAuthState::from_config(config)
+        .await
+        .expect("Failed to create auth state");
+
+    // Create a token WITHOUT kid in header
+    let claims = create_claims("user@example.com", vec!["admin"]);
+    let mut header = Header::new(jsonwebtoken::Algorithm::RS256);
+    header.kid = None; // Explicitly no kid
+
+    let key = EncodingKey::from_rsa_pem(TEST_KEYS.private_key_pem.as_bytes())
+        .expect("Failed to create encoding key");
+    let token_without_kid = encode(&header, &claims, &key).expect("Failed to encode token");
+
+    let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set");
+    let result = jwt_validator.validate(&token_without_kid).await;
+
+    assert!(result.is_err(), "Token without kid should fail validation");
+    let err = result.unwrap_err();
+    assert!(
+        err.to_string().contains("kid") || err.to_string().contains("Missing"),
+        "Error should mention missing kid: {}",
+        err
+    );
+}
+
+// ============================================================================
+// Edge Cases and Security Tests
+// ============================================================================
+
+#[tokio::test]
+async fn test_jwt_role_extraction_from_groups_claim() {
+    let (addr, _server) = start_mock_jwks_server().await;
+
+    // JWT validator checks alternate claims ("groups", "roles", etc.) as fallback
+    let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway")
+        .with_jwks_uri(format!(
+            "http://127.0.0.1:{}/.well-known/jwks.json",
+            addr.port()
+        ))
+        .with_role_mapping("administrators", Role::Admin);
+
+    let config = ControlPlaneAuthConfig {
+        jwt: Some(jwt_config),
+        api_keys: vec![],
+        audit_enabled: false,
+    };
+
+    let state = ControlPlaneAuthState::from_config(config)
+        .await
+        .expect("Failed to create auth state");
+
+    // Create token with 'groups' claim instead of 'roles'
+    let now = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_secs();
+
+    #[derive(Serialize)]
+    struct GroupsClaims {
+        sub: String,
+        iss: String,
+        aud: String,
+        exp: u64,
+        iat: u64,
+        jti: String,
+        groups: Vec<String>, // Use 'groups' instead of 'roles'
+    }
+
+    let claims = GroupsClaims {
+        sub: "group-user@example.com".to_string(),
+        iss: "https://test-issuer.example.com".to_string(),
+        aud: "test-gateway".to_string(),
+        exp: now + 3600,
+        iat: now,
+        jti: uuid::Uuid::new_v4().to_string(),
+        groups: vec!["administrators".to_string(), "developers".to_string()],
+    };
+
+    let mut header = Header::new(jsonwebtoken::Algorithm::RS256);
+    header.kid = Some(TEST_KEY_ID.to_string());
+    let key = EncodingKey::from_rsa_pem(TEST_KEYS.private_key_pem.as_bytes())
+        .expect("Failed to create encoding key");
+    let token = encode(&header, &claims, &key).expect("Failed to encode token");
+
+    let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set");
+    let result = jwt_validator.validate(&token).await;
+
+    assert!(
+        result.is_ok(),
+        "Token with groups claim should validate: {:?}",
+        result.err()
+    );
+    let validated = result.unwrap();
+    assert_eq!(
+        validated.role,
+        Role::Admin,
+        "Should map 'administrators' to Admin role"
+    );
+}
+
+#[tokio::test]
+async fn test_jwt_no_role_defaults_to_user() {
+    let (addr, _server) = start_mock_jwks_server().await;
+
+    let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway")
+        .with_jwks_uri(format!(
+            "http://127.0.0.1:{}/.well-known/jwks.json",
+            addr.port()
+        ));
+    // No role mapping configured
+
+    let config = ControlPlaneAuthConfig {
+        jwt: Some(jwt_config),
+        api_keys: vec![],
+        audit_enabled: false,
+    };
+
+    let state = ControlPlaneAuthState::from_config(config)
+        .await
+        .expect("Failed to create auth state");
+
+    // Create token without any role claims
+    let now = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_secs();
+
+    #[derive(Serialize)]
+    struct MinimalClaims {
+        sub: String,
+        iss: String,
+        aud: String,
+        exp: u64,
+        iat: u64,
+        jti: String,
+        // No roles, groups, or role claims
+    }
+
+    let claims = MinimalClaims {
+        sub: "minimal-user@example.com".to_string(),
+        iss: "https://test-issuer.example.com".to_string(),
+        aud: "test-gateway".to_string(),
+        exp: now + 3600,
+        iat: now,
+        jti: uuid::Uuid::new_v4().to_string(),
+    };
+
+    let mut header = Header::new(jsonwebtoken::Algorithm::RS256);
+    header.kid = Some(TEST_KEY_ID.to_string());
+    let key = EncodingKey::from_rsa_pem(TEST_KEYS.private_key_pem.as_bytes())
+        .expect("Failed to create encoding key");
+    let token = encode(&header, &claims, &key).expect("Failed to encode token");
+
+    let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set");
+    let result = jwt_validator.validate(&token).await;
+
+    assert!(result.is_ok(), "Token without role should still validate");
+    let validated = result.unwrap();
+    assert_eq!(
+        validated.role,
+        Role::User,
+        "Should default to User role when no role found"
+    );
+}
+
+#[tokio::test]
+async fn test_multiple_api_keys() {
+    let config = ControlPlaneAuthConfig {
+        jwt: None,
+        api_keys: vec![
+            ApiKeyEntry::new("admin-1", "Admin Key 1", "sk-admin-key-1", Role::Admin),
+            ApiKeyEntry::new("admin-2", "Admin Key 2", "sk-admin-key-2", Role::Admin),
+            ApiKeyEntry::new("user-1", "User Key 1", "sk-user-key-1", Role::User),
+        ],
+        audit_enabled: false,
+    };
+
+    // All keys should be findable
+    assert!(config.find_api_key("sk-admin-key-1").is_some());
+    assert!(config.find_api_key("sk-admin-key-2").is_some());
+    assert!(config.find_api_key("sk-user-key-1").is_some());
+
+    // Verify correct roles
+    assert_eq!(
+        config.find_api_key("sk-admin-key-1").unwrap().role,
+        Role::Admin
+    );
+    assert_eq!(
+        config.find_api_key("sk-user-key-1").unwrap().role,
+        Role::User
+    );
+
+    // Non-existent key should not be found
+    assert!(config.find_api_key("sk-nonexistent").is_none());
+}
+
+#[tokio::test]
+async fn test_config_is_enabled_checks() {
+    // Empty config is not enabled
+    let empty_config = ControlPlaneAuthConfig {
+        jwt: None,
+        api_keys: vec![],
+        audit_enabled: false,
+    };
+    assert!(!empty_config.is_enabled());
+
+    // Config with only API keys is enabled
+    let api_key_config = ControlPlaneAuthConfig {
+        jwt: None,
+        api_keys: vec![ApiKeyEntry::new("key-1", "Test", "sk-test", Role::Admin)],
+        audit_enabled: false,
+    };
+    assert!(api_key_config.is_enabled());
+
+    // Config with only JWT is enabled
+    let jwt_config = ControlPlaneAuthConfig {
+        jwt: Some(JwtConfig::new("https://issuer.example.com", "audience")),
+        api_keys: vec![],
+        audit_enabled: false,
+    };
+    assert!(jwt_config.is_enabled());
+
+    // Config with both is enabled
+    let full_config = ControlPlaneAuthConfig {
+        jwt: Some(JwtConfig::new("https://issuer.example.com", "audience")),
+        api_keys: vec![ApiKeyEntry::new("key-1", "Test", "sk-test", Role::Admin)],
+        audit_enabled: true,
+    };
+    assert!(full_config.is_enabled());
+}

third_party/sglang/sgl-model-gateway/tests/security/auth_test.rs

@@ -0,0 +1,243 @@
+//! Authentication and authorization integration tests
+//!
+//! Tests for API key enforcement and access control.
+
+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 AUTH_HEADER: &str = "Authorization";
+
+#[cfg(test)]
+mod auth_tests {
+    use super::*;
+
+    /// Test request without API key when auth is not required
+    #[tokio::test]
+    async fn test_no_auth_required() {
+        let config = TestRouterConfig::round_robin(4300);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20300)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Request without auth header should succeed when no auth required
+        let payload = json!({
+            "text": "Test without auth",
+            "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 without auth should succeed when no auth required"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test request with valid API key format
+    #[tokio::test]
+    async fn test_with_api_key_header() {
+        let config = TestRouterConfig::round_robin(4301);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20301)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Request with Bearer token header
+        let payload = json!({
+            "text": "Test with auth header",
+            "stream": false
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/generate")
+            .header(CONTENT_TYPE, "application/json")
+            .header(AUTH_HEADER, "Bearer test-api-key-12345")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        // Without auth enforcement, request should succeed
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Request with auth header should succeed"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test health endpoint doesn't require authentication
+    #[tokio::test]
+    async fn test_health_endpoint_no_auth() {
+        let config = TestRouterConfig::round_robin(4302);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20302)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Health endpoint should be accessible without auth
+        let req = Request::builder()
+            .method("GET")
+            .uri("/health")
+            .body(Body::empty())
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Health endpoint should not require auth"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test OpenAI-compatible API key header (X-API-Key)
+    #[tokio::test]
+    async fn test_openai_api_key_header() {
+        let config = TestRouterConfig::round_robin(4303);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20303)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Request with X-API-Key header (OpenAI style)
+        let payload = json!({
+            "text": "Test with X-API-Key",
+            "stream": false
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/generate")
+            .header(CONTENT_TYPE, "application/json")
+            .header("X-API-Key", "sk-test-key-12345")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Request with X-API-Key should succeed"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test multiple concurrent authenticated requests
+    #[tokio::test]
+    async fn test_concurrent_authenticated_requests() {
+        use std::sync::{
+            atomic::{AtomicUsize, Ordering},
+            Arc,
+        };
+
+        let config = TestRouterConfig::round_robin(4304);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20304)]).await;
+
+        let app = ctx.create_app().await;
+        let success_count = Arc::new(AtomicUsize::new(0));
+        let mut handles = Vec::new();
+
+        for i in 0..20 {
+            let app_clone = app.clone();
+            let success_clone = Arc::clone(&success_count);
+
+            let handle = tokio::spawn(async move {
+                let payload = json!({
+                    "text": format!("Concurrent auth test {}", i),
+                    "stream": false
+                });
+
+                let req = Request::builder()
+                    .method("POST")
+                    .uri("/generate")
+                    .header(CONTENT_TYPE, "application/json")
+                    .header(AUTH_HEADER, format!("Bearer test-key-{}", i))
+                    .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();
+        }
+
+        assert_eq!(
+            success_count.load(Ordering::SeqCst),
+            20,
+            "All concurrent authenticated requests should succeed"
+        );
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod mtls_tests {
+    use super::*;
+
+    /// Test that TLS configuration options exist
+    /// Note: Actual mTLS testing would require certificate setup
+    #[tokio::test]
+    async fn test_tls_config_available() {
+        // This test verifies the config builder accepts TLS-related options
+        // Actual mTLS testing requires certificate infrastructure
+        let config = TestRouterConfig::round_robin(4305);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20305)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Basic request should work (no TLS in test mode)
+        let payload = json!({
+            "text": "TLS config test",
+            "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);
+
+        ctx.shutdown().await;
+    }
+}

third_party/sglang/sgl-model-gateway/tests/security/mod.rs

@@ -0,0 +1,5 @@
+//! Security integration tests
+
+pub mod auth_integration_test;
+pub mod auth_test;
+pub mod mtls_test;

third_party/sglang/sgl-model-gateway/tests/security/mtls_test.rs

@@ -0,0 +1,506 @@
+//! mTLS (Mutual TLS) integration tests
+//!
+//! Tests for TLS and mTLS communication between router and workers.
+//! Covers:
+//! - Successful mTLS communication with client certificates
+//! - TLS failure without client certificate when required
+//! - TLS-only mode (server authentication only)
+//! - TLS failure without CA certificate
+
+use std::{io::BufReader, time::Duration};
+
+use rustls::{ClientConfig, RootCertStore};
+use rustls_pemfile::{certs, pkcs8_private_keys};
+use serde_json::json;
+
+use crate::common::{
+    test_certs::TestCertificates,
+    tls_mock_worker::{TlsMockWorker, TlsMockWorkerConfig},
+};
+
+// Ensure crypto provider is installed
+fn ensure_crypto_provider() {
+    use std::sync::Once;
+    static INIT: Once = Once::new();
+    INIT.call_once(|| {
+        let _ = rustls::crypto::ring::default_provider().install_default();
+    });
+}
+
+/// Helper to create a TLS-enabled reqwest client with client certificates
+fn create_mtls_client(
+    test_certs: &TestCertificates,
+) -> Result<reqwest::Client, Box<dyn std::error::Error>> {
+    ensure_crypto_provider();
+
+    // Read CA certificate
+    let ca_file = std::fs::File::open(&test_certs.ca_cert_path)?;
+    let mut ca_reader = BufReader::new(ca_file);
+    let ca_certs: Vec<_> = certs(&mut ca_reader).filter_map(|r| r.ok()).collect();
+
+    let mut root_store = RootCertStore::empty();
+    for cert in ca_certs {
+        root_store.add(cert)?;
+    }
+
+    // Read client certificate
+    let cert_file = std::fs::File::open(&test_certs.client_cert_path)?;
+    let mut cert_reader = BufReader::new(cert_file);
+    let client_certs: Vec<_> = certs(&mut cert_reader).filter_map(|r| r.ok()).collect();
+
+    // Read client key
+    let key_file = std::fs::File::open(&test_certs.client_key_path)?;
+    let mut key_reader = BufReader::new(key_file);
+    let client_key = pkcs8_private_keys(&mut key_reader)
+        .next()
+        .ok_or("No private key found")??;
+
+    // Build client config with client certificate
+    let client_config = ClientConfig::builder()
+        .with_root_certificates(root_store)
+        .with_client_auth_cert(client_certs, client_key.into())
+        .map_err(|e| format!("Failed to build client config: {}", e))?;
+
+    let client = reqwest::Client::builder()
+        .use_preconfigured_tls(client_config)
+        .timeout(Duration::from_secs(10))
+        .build()?;
+
+    Ok(client)
+}
+
+/// Helper to create a TLS client without client certificates (for server-auth-only)
+fn create_tls_client(
+    test_certs: &TestCertificates,
+) -> Result<reqwest::Client, Box<dyn std::error::Error>> {
+    ensure_crypto_provider();
+
+    // Read CA certificate
+    let ca_file = std::fs::File::open(&test_certs.ca_cert_path)?;
+    let mut ca_reader = BufReader::new(ca_file);
+    let ca_certs: Vec<_> = certs(&mut ca_reader).filter_map(|r| r.ok()).collect();
+
+    let mut root_store = RootCertStore::empty();
+    for cert in ca_certs {
+        root_store.add(cert)?;
+    }
+
+    // Build client config without client certificate
+    let client_config = ClientConfig::builder()
+        .with_root_certificates(root_store)
+        .with_no_client_auth();
+
+    let client = reqwest::Client::builder()
+        .use_preconfigured_tls(client_config)
+        .timeout(Duration::from_secs(10))
+        .build()?;
+
+    Ok(client)
+}
+
+/// Helper to create a client without CA certificate (for failure test)
+fn create_client_without_ca() -> Result<reqwest::Client, Box<dyn std::error::Error>> {
+    ensure_crypto_provider();
+
+    // Build client with empty root store (will fail to verify server)
+    let root_store = RootCertStore::empty();
+    let client_config = ClientConfig::builder()
+        .with_root_certificates(root_store)
+        .with_no_client_auth();
+
+    let client = reqwest::Client::builder()
+        .use_preconfigured_tls(client_config)
+        .timeout(Duration::from_secs(10))
+        .build()?;
+
+    Ok(client)
+}
+
+#[cfg(test)]
+mod mtls_tests {
+    use super::*;
+
+    /// Test successful mTLS communication between client and TLS worker
+    ///
+    /// This test verifies that:
+    /// 1. TLS mock worker starts with mTLS configuration
+    /// 2. Client with proper certificates can connect and communicate
+    /// 3. Requests succeed with proper authentication
+    #[tokio::test]
+    async fn test_mtls_successful_communication() {
+        // Generate test certificates
+        let certs = TestCertificates::generate().expect("Failed to generate test certificates");
+
+        // Start mTLS-enabled mock worker
+        let mut worker = TlsMockWorker::new(TlsMockWorkerConfig {
+            port: 0, // Auto-assign port
+            require_client_cert: true,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        });
+
+        let worker_url = worker
+            .start(
+                &certs.server_cert_path,
+                &certs.server_key_path,
+                Some(&certs.ca_cert_path),
+            )
+            .await
+            .expect("Failed to start mTLS worker");
+
+        // Create client with full mTLS credentials
+        let client = create_mtls_client(&certs).expect("Failed to create mTLS client");
+
+        // Test health endpoint
+        let health_resp = client
+            .get(format!("{}/health", worker_url))
+            .send()
+            .await
+            .expect("Health request failed");
+
+        assert!(
+            health_resp.status().is_success(),
+            "Health check should succeed with mTLS: {}",
+            health_resp.status()
+        );
+
+        let health_json: serde_json::Value = health_resp.json().await.unwrap();
+        assert_eq!(health_json["status"], "healthy");
+        assert_eq!(health_json["tls_enabled"], true);
+
+        // Test generate endpoint
+        let payload = json!({
+            "text": "Test mTLS request",
+            "stream": false
+        });
+
+        let gen_resp = client
+            .post(format!("{}/generate", worker_url))
+            .json(&payload)
+            .send()
+            .await
+            .expect("Generate request failed");
+
+        assert!(
+            gen_resp.status().is_success(),
+            "Generate should succeed with mTLS: {}",
+            gen_resp.status()
+        );
+
+        let gen_json: serde_json::Value = gen_resp.json().await.unwrap();
+        assert!(gen_json["text"].as_str().is_some());
+        assert_eq!(gen_json["meta_info"]["tls_verified"], true);
+
+        // Cleanup
+        worker.stop().await;
+    }
+
+    /// Test that mTLS worker rejects connections without client certificate
+    ///
+    /// This test verifies that:
+    /// 1. mTLS worker requires client certificate
+    /// 2. Connection without client cert fails
+    #[tokio::test]
+    async fn test_mtls_failure_without_client_cert() {
+        // Generate test certificates
+        let certs = TestCertificates::generate().expect("Failed to generate test certificates");
+
+        // Start mTLS-enabled mock worker (requires client cert)
+        let mut worker = TlsMockWorker::new(TlsMockWorkerConfig {
+            port: 0,
+            require_client_cert: true,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        });
+
+        let worker_url = worker
+            .start(
+                &certs.server_cert_path,
+                &certs.server_key_path,
+                Some(&certs.ca_cert_path),
+            )
+            .await
+            .expect("Failed to start mTLS worker");
+
+        // Create client WITHOUT client certificate (only has CA for server verification)
+        let client = create_tls_client(&certs).expect("Failed to create TLS-only client");
+
+        // Attempt to connect - should fail because no client cert provided
+        let result = client.get(format!("{}/health", worker_url)).send().await;
+
+        // Connection should fail or be rejected
+        assert!(
+            result.is_err(),
+            "Connection should fail without client certificate"
+        );
+
+        // Cleanup
+        worker.stop().await;
+    }
+
+    /// Test TLS-only mode (server authentication only, no client cert required)
+    ///
+    /// This test verifies that:
+    /// 1. TLS worker can operate without requiring client certificates
+    /// 2. Client can connect with just CA certificate for server verification
+    #[tokio::test]
+    async fn test_tls_server_auth_only() {
+        // Generate test certificates
+        let certs = TestCertificates::generate().expect("Failed to generate test certificates");
+
+        // Start TLS-only mock worker (does NOT require client cert)
+        let mut worker = TlsMockWorker::new(TlsMockWorkerConfig {
+            port: 0,
+            require_client_cert: false, // TLS-only mode
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        });
+
+        let worker_url = worker
+            .start(
+                &certs.server_cert_path,
+                &certs.server_key_path,
+                None, // No CA needed for client verification
+            )
+            .await
+            .expect("Failed to start TLS worker");
+
+        // Create client with just CA cert for server verification (no client cert)
+        let client = create_tls_client(&certs).expect("Failed to create TLS client");
+
+        // Test health endpoint - should succeed without client cert
+        let health_resp = client
+            .get(format!("{}/health", worker_url))
+            .send()
+            .await
+            .expect("Health request failed");
+
+        assert!(
+            health_resp.status().is_success(),
+            "Health check should succeed in TLS-only mode: {}",
+            health_resp.status()
+        );
+
+        let health_json: serde_json::Value = health_resp.json().await.unwrap();
+        assert_eq!(health_json["status"], "healthy");
+
+        // Test chat completions endpoint
+        let payload = json!({
+            "model": "mock-tls-model",
+            "messages": [{"role": "user", "content": "Hello TLS"}]
+        });
+
+        let chat_resp = client
+            .post(format!("{}/v1/chat/completions", worker_url))
+            .json(&payload)
+            .send()
+            .await
+            .expect("Chat request failed");
+
+        assert!(
+            chat_resp.status().is_success(),
+            "Chat should succeed in TLS-only mode: {}",
+            chat_resp.status()
+        );
+
+        // Cleanup
+        worker.stop().await;
+    }
+
+    /// Test TLS failure when client doesn't have CA certificate
+    ///
+    /// This test verifies that:
+    /// 1. Client cannot verify server without proper CA certificate
+    /// 2. Connection fails due to certificate verification
+    #[tokio::test]
+    async fn test_tls_failure_without_ca_cert() {
+        // Generate test certificates
+        let certs = TestCertificates::generate().expect("Failed to generate test certificates");
+
+        // Start TLS mock worker
+        let mut worker = TlsMockWorker::new(TlsMockWorkerConfig {
+            port: 0,
+            require_client_cert: false,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        });
+
+        let worker_url = worker
+            .start(&certs.server_cert_path, &certs.server_key_path, None)
+            .await
+            .expect("Failed to start TLS worker");
+
+        // Create client WITHOUT CA certificate
+        let client = create_client_without_ca().expect("Failed to create client without CA");
+
+        // Attempt to connect - should fail because cannot verify server cert
+        let result = client.get(format!("{}/health", worker_url)).send().await;
+
+        // Connection should fail due to certificate verification
+        assert!(
+            result.is_err(),
+            "Connection should fail without CA certificate for server verification"
+        );
+
+        // Cleanup
+        worker.stop().await;
+    }
+
+    /// Test multiple concurrent mTLS requests
+    ///
+    /// This test verifies that mTLS connections work correctly under concurrent load
+    #[tokio::test]
+    async fn test_mtls_concurrent_requests() {
+        use std::sync::{
+            atomic::{AtomicUsize, Ordering},
+            Arc,
+        };
+
+        // Generate test certificates
+        let certs = TestCertificates::generate().expect("Failed to generate test certificates");
+
+        // Start mTLS-enabled mock worker
+        let mut worker = TlsMockWorker::new(TlsMockWorkerConfig {
+            port: 0,
+            require_client_cert: true,
+            response_delay_ms: 10, // Small delay to simulate work
+            fail_rate: 0.0,
+        });
+
+        let worker_url = worker
+            .start(
+                &certs.server_cert_path,
+                &certs.server_key_path,
+                Some(&certs.ca_cert_path),
+            )
+            .await
+            .expect("Failed to start mTLS worker");
+
+        // Create mTLS client
+        let client = Arc::new(create_mtls_client(&certs).expect("Failed to create mTLS client"));
+        let success_count = Arc::new(AtomicUsize::new(0));
+        let worker_url = Arc::new(worker_url);
+
+        let mut handles = Vec::new();
+
+        // Spawn concurrent requests
+        for i in 0..10 {
+            let client_clone = Arc::clone(&client);
+            let success_clone = Arc::clone(&success_count);
+            let url_clone = Arc::clone(&worker_url);
+
+            let handle = tokio::spawn(async move {
+                let payload = json!({
+                    "text": format!("Concurrent mTLS request {}", i),
+                    "stream": false
+                });
+
+                let resp = client_clone
+                    .post(format!("{}/generate", url_clone))
+                    .json(&payload)
+                    .send()
+                    .await;
+
+                if let Ok(response) = resp {
+                    if response.status().is_success() {
+                        success_clone.fetch_add(1, Ordering::SeqCst);
+                    }
+                }
+            });
+
+            handles.push(handle);
+        }
+
+        // Wait for all requests to complete
+        for handle in handles {
+            handle.await.unwrap();
+        }
+
+        // All requests should succeed
+        assert_eq!(
+            success_count.load(Ordering::SeqCst),
+            10,
+            "All concurrent mTLS requests should succeed"
+        );
+
+        // Cleanup
+        worker.stop().await;
+    }
+}
+
+#[cfg(test)]
+mod certificate_tests {
+    use super::*;
+
+    /// Test that certificate generation works correctly
+    #[test]
+    fn test_certificate_generation() {
+        let certs = TestCertificates::generate().expect("Failed to generate certificates");
+
+        // Verify all files exist and are not empty
+        assert!(certs.ca_cert_path.exists());
+        assert!(certs.ca_key_path.exists());
+        assert!(certs.server_cert_path.exists());
+        assert!(certs.server_key_path.exists());
+        assert!(certs.client_cert_path.exists());
+        assert!(certs.client_key_path.exists());
+
+        // Verify PEM format
+        let ca_cert_content = std::fs::read_to_string(&certs.ca_cert_path).unwrap();
+        assert!(ca_cert_content.contains("-----BEGIN CERTIFICATE-----"));
+        assert!(ca_cert_content.contains("-----END CERTIFICATE-----"));
+
+        let server_key_content = std::fs::read_to_string(&certs.server_key_path).unwrap();
+        assert!(server_key_content.contains("-----BEGIN PRIVATE KEY-----"));
+        assert!(server_key_content.contains("-----END PRIVATE KEY-----"));
+    }
+
+    /// Test that generated certificates can be parsed by rustls
+    #[test]
+    fn test_certificate_parsing() {
+        ensure_crypto_provider();
+
+        let test_certs = TestCertificates::generate().expect("Failed to generate certificates");
+
+        // Verify CA certificate can be parsed by rustls
+        let ca_file = std::fs::File::open(&test_certs.ca_cert_path).unwrap();
+        let mut ca_reader = BufReader::new(ca_file);
+        let ca_certs: Vec<_> = certs(&mut ca_reader).filter_map(|r| r.ok()).collect();
+        assert!(!ca_certs.is_empty(), "CA certificate should be parseable");
+
+        // Verify client certificate can be parsed
+        let cert_file = std::fs::File::open(&test_certs.client_cert_path).unwrap();
+        let mut cert_reader = BufReader::new(cert_file);
+        let client_certs: Vec<_> = certs(&mut cert_reader).filter_map(|r| r.ok()).collect();
+        assert!(
+            !client_certs.is_empty(),
+            "Client certificate should be parseable"
+        );
+
+        // Verify client key can be parsed
+        let key_file = std::fs::File::open(&test_certs.client_key_path).unwrap();
+        let mut key_reader = BufReader::new(key_file);
+        let keys: Vec<_> = pkcs8_private_keys(&mut key_reader)
+            .filter_map(|r| r.ok())
+            .collect();
+        assert!(!keys.is_empty(), "Client key should be parseable");
+
+        // Verify we can build a complete client config
+        let mut root_store = RootCertStore::empty();
+        for cert in ca_certs {
+            root_store
+                .add(cert)
+                .expect("Should be able to add CA cert to store");
+        }
+
+        let key = keys.into_iter().next().unwrap();
+        let config_result = ClientConfig::builder()
+            .with_root_certificates(root_store)
+            .with_client_auth_cert(client_certs, key.into());
+
+        assert!(
+            config_result.is_ok(),
+            "Should be able to create client config with certs"
+        );
+    }
+}

third_party/sglang/sgl-model-gateway/tests/security_tests.rs

@@ -0,0 +1,8 @@
+//! Security integration tests
+
+#[path = "common/mod.rs"]
+pub mod common;
+
+mod security;
+
+pub use security::*;

third_party/sglang/sgl-model-gateway/tests/spec/chat_completion.rs

@@ -0,0 +1,570 @@
+use serde_json::json;
+use smg::protocols::{
+    chat::{ChatCompletionRequest, ChatMessage, MessageContent},
+    common::{
+        Function, FunctionCall, FunctionChoice, StreamOptions, Tool, ToolChoice, ToolChoiceValue,
+        ToolReference,
+    },
+    validated::Normalizable,
+};
+use validator::Validate;
+
+// Deprecated fields normalization tests
+
+#[test]
+fn test_max_tokens_normalizes_to_max_completion_tokens() {
+    #[allow(deprecated)]
+    let mut req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        max_tokens: Some(100),
+        max_completion_tokens: None,
+        ..Default::default()
+    };
+
+    req.normalize();
+    assert_eq!(
+        req.max_completion_tokens,
+        Some(100),
+        "max_tokens should be copied to max_completion_tokens"
+    );
+    #[allow(deprecated)]
+    {
+        assert!(
+            req.max_tokens.is_none(),
+            "Deprecated field should be cleared"
+        );
+    }
+    assert!(
+        req.validate().is_ok(),
+        "Should be valid after normalization"
+    );
+}
+
+#[test]
+fn test_max_completion_tokens_takes_precedence() {
+    #[allow(deprecated)]
+    let mut req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        max_tokens: Some(100),
+        max_completion_tokens: Some(200),
+        ..Default::default()
+    };
+
+    req.normalize();
+    assert_eq!(
+        req.max_completion_tokens,
+        Some(200),
+        "max_completion_tokens should take precedence"
+    );
+    assert!(
+        req.validate().is_ok(),
+        "Should be valid after normalization"
+    );
+}
+
+#[test]
+fn test_functions_normalizes_to_tools() {
+    #[allow(deprecated)]
+    let mut req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        functions: Some(vec![Function {
+            name: "test_func".to_string(),
+            description: Some("Test function".to_string()),
+            parameters: json!({}),
+            strict: None,
+        }]),
+        tools: None,
+        ..Default::default()
+    };
+
+    req.normalize();
+    assert!(req.tools.is_some(), "functions should be migrated to tools");
+    assert_eq!(req.tools.as_ref().unwrap().len(), 1);
+    assert_eq!(req.tools.as_ref().unwrap()[0].function.name, "test_func");
+    #[allow(deprecated)]
+    {
+        assert!(
+            req.functions.is_none(),
+            "Deprecated field should be cleared"
+        );
+    }
+    assert!(
+        req.validate().is_ok(),
+        "Should be valid after normalization"
+    );
+}
+
+#[test]
+fn test_function_call_normalizes_to_tool_choice() {
+    #[allow(deprecated)]
+    let mut req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        function_call: Some(FunctionCall::None),
+        tool_choice: None,
+        ..Default::default()
+    };
+
+    req.normalize();
+    assert!(
+        req.tool_choice.is_some(),
+        "function_call should be migrated to tool_choice"
+    );
+    assert!(matches!(
+        req.tool_choice,
+        Some(ToolChoice::Value(ToolChoiceValue::None))
+    ));
+    #[allow(deprecated)]
+    {
+        assert!(
+            req.function_call.is_none(),
+            "Deprecated field should be cleared"
+        );
+    }
+    assert!(
+        req.validate().is_ok(),
+        "Should be valid after normalization"
+    );
+}
+
+#[test]
+fn test_function_call_function_variant_normalizes() {
+    #[allow(deprecated)]
+    let mut req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        function_call: Some(FunctionCall::Function {
+            name: "my_function".to_string(),
+        }),
+        tool_choice: None,
+        tools: Some(vec![Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "my_function".to_string(),
+                description: None,
+                parameters: json!({}),
+                strict: None,
+            },
+        }]),
+        ..Default::default()
+    };
+
+    req.normalize();
+    assert!(
+        req.tool_choice.is_some(),
+        "function_call should be migrated to tool_choice"
+    );
+    match &req.tool_choice {
+        Some(ToolChoice::Function { function, .. }) => {
+            assert_eq!(function.name, "my_function");
+        }
+        _ => panic!("Expected ToolChoice::Function variant"),
+    }
+    #[allow(deprecated)]
+    {
+        assert!(
+            req.function_call.is_none(),
+            "Deprecated field should be cleared"
+        );
+    }
+    assert!(
+        req.validate().is_ok(),
+        "Should be valid after normalization"
+    );
+}
+
+// Stream options validation tests
+
+#[test]
+fn test_stream_options_requires_stream_enabled() {
+    let req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        stream: false,
+        stream_options: Some(StreamOptions {
+            include_usage: Some(true),
+        }),
+        ..Default::default()
+    };
+
+    let result = req.validate();
+    assert!(
+        result.is_err(),
+        "Should reject stream_options when stream is false"
+    );
+    let err = result.unwrap_err().to_string();
+    assert!(
+        err.contains("stream_options") && err.contains("stream") && err.contains("enabled"),
+        "Error should mention stream dependency: {}",
+        err
+    );
+}
+
+#[test]
+fn test_stream_options_valid_when_stream_enabled() {
+    let req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        stream: true,
+        stream_options: Some(StreamOptions {
+            include_usage: Some(true),
+        }),
+        ..Default::default()
+    };
+
+    let result = req.validate();
+    assert!(
+        result.is_ok(),
+        "Should accept stream_options when stream is true"
+    );
+}
+
+#[test]
+fn test_no_stream_options_valid_when_stream_disabled() {
+    let req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        stream: false,
+        stream_options: None,
+        ..Default::default()
+    };
+
+    let result = req.validate();
+    assert!(
+        result.is_ok(),
+        "Should accept no stream_options when stream is false"
+    );
+}
+
+// Tool choice validation tests
+#[test]
+fn test_tool_choice_function_not_found() {
+    let req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        tools: Some(vec![Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "get_weather".to_string(),
+                description: Some("Get weather".to_string()),
+                parameters: json!({}),
+                strict: None,
+            },
+        }]),
+        tool_choice: Some(ToolChoice::Function {
+            function: FunctionChoice {
+                name: "nonexistent_function".to_string(),
+            },
+            tool_type: "function".to_string(),
+        }),
+        ..Default::default()
+    };
+
+    let result = req.validate();
+    assert!(result.is_err(), "Should reject nonexistent function name");
+    let err = result.unwrap_err().to_string();
+    assert!(
+        err.contains("function 'nonexistent_function' not found"),
+        "Error should mention the missing function: {}",
+        err
+    );
+}
+
+#[test]
+fn test_tool_choice_function_exists_valid() {
+    let req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        tools: Some(vec![Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "get_weather".to_string(),
+                description: Some("Get weather".to_string()),
+                parameters: json!({}),
+                strict: None,
+            },
+        }]),
+        tool_choice: Some(ToolChoice::Function {
+            function: FunctionChoice {
+                name: "get_weather".to_string(),
+            },
+            tool_type: "function".to_string(),
+        }),
+        ..Default::default()
+    };
+
+    let result = req.validate();
+    assert!(result.is_ok(), "Should accept existing function name");
+}
+
+#[test]
+fn test_tool_choice_allowed_tools_invalid_mode() {
+    let req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        tools: Some(vec![Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "get_weather".to_string(),
+                description: Some("Get weather".to_string()),
+                parameters: json!({}),
+                strict: None,
+            },
+        }]),
+        tool_choice: Some(ToolChoice::AllowedTools {
+            mode: "invalid_mode".to_string(),
+            tools: vec![ToolReference::Function {
+                name: "get_weather".to_string(),
+            }],
+            tool_type: "function".to_string(),
+        }),
+        ..Default::default()
+    };
+
+    let result = req.validate();
+    assert!(result.is_err(), "Should reject invalid mode");
+    let err = result.unwrap_err().to_string();
+    assert!(
+        err.contains("must be 'auto' or 'required'"),
+        "Error should mention valid modes: {}",
+        err
+    );
+}
+
+#[test]
+fn test_tool_choice_allowed_tools_valid_mode_auto() {
+    let req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        tools: Some(vec![Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "get_weather".to_string(),
+                description: Some("Get weather".to_string()),
+                parameters: json!({}),
+                strict: None,
+            },
+        }]),
+        tool_choice: Some(ToolChoice::AllowedTools {
+            mode: "auto".to_string(),
+            tools: vec![ToolReference::Function {
+                name: "get_weather".to_string(),
+            }],
+            tool_type: "function".to_string(),
+        }),
+        ..Default::default()
+    };
+
+    let result = req.validate();
+    assert!(result.is_ok(), "Should accept 'auto' mode");
+}
+
+#[test]
+fn test_tool_choice_allowed_tools_valid_mode_required() {
+    let req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        tools: Some(vec![Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "get_weather".to_string(),
+                description: Some("Get weather".to_string()),
+                parameters: json!({}),
+                strict: None,
+            },
+        }]),
+        tool_choice: Some(ToolChoice::AllowedTools {
+            mode: "required".to_string(),
+            tools: vec![ToolReference::Function {
+                name: "get_weather".to_string(),
+            }],
+            tool_type: "function".to_string(),
+        }),
+        ..Default::default()
+    };
+
+    let result = req.validate();
+    assert!(result.is_ok(), "Should accept 'required' mode");
+}
+
+#[test]
+fn test_tool_choice_allowed_tools_tool_not_found() {
+    let req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        tools: Some(vec![Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "get_weather".to_string(),
+                description: Some("Get weather".to_string()),
+                parameters: json!({}),
+                strict: None,
+            },
+        }]),
+        tool_choice: Some(ToolChoice::AllowedTools {
+            mode: "auto".to_string(),
+            tools: vec![ToolReference::Function {
+                name: "nonexistent_tool".to_string(),
+            }],
+            tool_type: "function".to_string(),
+        }),
+        ..Default::default()
+    };
+
+    let result = req.validate();
+    assert!(result.is_err(), "Should reject nonexistent tool name");
+    let err = result.unwrap_err().to_string();
+    assert!(
+        err.contains("tool 'nonexistent_tool' not found"),
+        "Error should mention the missing tool: {}",
+        err
+    );
+}
+
+#[test]
+fn test_tool_choice_allowed_tools_multiple_tools_valid() {
+    let req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        tools: Some(vec![
+            Tool {
+                tool_type: "function".to_string(),
+                function: Function {
+                    name: "get_weather".to_string(),
+                    description: Some("Get weather".to_string()),
+                    parameters: json!({}),
+                    strict: None,
+                },
+            },
+            Tool {
+                tool_type: "function".to_string(),
+                function: Function {
+                    name: "get_time".to_string(),
+                    description: Some("Get time".to_string()),
+                    parameters: json!({}),
+                    strict: None,
+                },
+            },
+        ]),
+        tool_choice: Some(ToolChoice::AllowedTools {
+            mode: "auto".to_string(),
+            tools: vec![
+                ToolReference::Function {
+                    name: "get_weather".to_string(),
+                },
+                ToolReference::Function {
+                    name: "get_time".to_string(),
+                },
+            ],
+            tool_type: "function".to_string(),
+        }),
+        ..Default::default()
+    };
+
+    let result = req.validate();
+    assert!(result.is_ok(), "Should accept all valid tool references");
+}
+
+#[test]
+fn test_tool_choice_allowed_tools_one_invalid_among_valid() {
+    let req = ChatCompletionRequest {
+        model: "test-model".to_string(),
+        messages: vec![ChatMessage::User {
+            content: MessageContent::Text("hello".to_string()),
+            name: None,
+        }],
+        tools: Some(vec![
+            Tool {
+                tool_type: "function".to_string(),
+                function: Function {
+                    name: "get_weather".to_string(),
+                    description: Some("Get weather".to_string()),
+                    parameters: json!({}),
+                    strict: None,
+                },
+            },
+            Tool {
+                tool_type: "function".to_string(),
+                function: Function {
+                    name: "get_time".to_string(),
+                    description: Some("Get time".to_string()),
+                    parameters: json!({}),
+                    strict: None,
+                },
+            },
+        ]),
+        tool_choice: Some(ToolChoice::AllowedTools {
+            mode: "auto".to_string(),
+            tools: vec![
+                ToolReference::Function {
+                    name: "get_weather".to_string(),
+                },
+                ToolReference::Function {
+                    name: "nonexistent_tool".to_string(),
+                },
+            ],
+            tool_type: "function".to_string(),
+        }),
+        ..Default::default()
+    };
+
+    let result = req.validate();
+    assert!(
+        result.is_err(),
+        "Should reject if any tool reference is invalid"
+    );
+    let err = result.unwrap_err().to_string();
+    assert!(
+        err.contains("tool 'nonexistent_tool' not found"),
+        "Error should mention the missing tool: {}",
+        err
+    );
+}

third_party/sglang/sgl-model-gateway/tests/spec/chat_message.rs

@@ -0,0 +1,91 @@
+use serde_json::json;
+use smg::protocols::chat::{ChatMessage, MessageContent};
+
+#[test]
+fn test_chat_message_tagged_by_role_system() {
+    let json = json!({
+        "role": "system",
+        "content": "You are a helpful assistant"
+    });
+
+    let msg: ChatMessage = serde_json::from_value(json).unwrap();
+    match msg {
+        ChatMessage::System { content, .. } => {
+            assert_eq!(
+                content,
+                MessageContent::Text("You are a helpful assistant".to_string())
+            )
+        }
+        _ => panic!("Expected System variant"),
+    }
+}
+
+#[test]
+fn test_chat_message_tagged_by_role_user() {
+    let json = json!({
+        "role": "user",
+        "content": "Hello"
+    });
+
+    let msg: ChatMessage = serde_json::from_value(json).unwrap();
+    match msg {
+        ChatMessage::User { content, .. } => match content {
+            MessageContent::Text(text) => assert_eq!(text, "Hello"),
+            _ => panic!("Expected text content"),
+        },
+        _ => panic!("Expected User variant"),
+    }
+}
+
+#[test]
+fn test_chat_message_tagged_by_role_assistant() {
+    let json = json!({
+        "role": "assistant",
+        "content": "Hi there!"
+    });
+
+    let msg: ChatMessage = serde_json::from_value(json).unwrap();
+    match msg {
+        ChatMessage::Assistant { content, .. } => {
+            assert_eq!(content, Some(MessageContent::Text("Hi there!".to_string())));
+        }
+        _ => panic!("Expected Assistant variant"),
+    }
+}
+
+#[test]
+fn test_chat_message_tagged_by_role_tool() {
+    let json = json!({
+        "role": "tool",
+        "content": "Tool result",
+        "tool_call_id": "call_123"
+    });
+
+    let msg: ChatMessage = serde_json::from_value(json).unwrap();
+    match msg {
+        ChatMessage::Tool {
+            content,
+            tool_call_id,
+        } => {
+            match content {
+                MessageContent::Text(text) => {
+                    assert_eq!(text, "Tool result");
+                }
+                _ => panic!("Expected content to be a string"),
+            }
+            assert_eq!(tool_call_id, "call_123");
+        }
+        _ => panic!("Expected Tool variant"),
+    }
+}
+
+#[test]
+fn test_chat_message_wrong_role_rejected() {
+    let json = json!({
+        "role": "invalid_role",
+        "content": "test"
+    });
+
+    let result = serde_json::from_value::<ChatMessage>(json);
+    assert!(result.is_err(), "Should reject invalid role");
+}

third_party/sglang/sgl-model-gateway/tests/spec/embedding.rs

@@ -0,0 +1,102 @@
+use serde_json::{from_str, json, to_string};
+use smg::protocols::{common::GenerationRequest, embedding::EmbeddingRequest};
+
+#[test]
+fn test_embedding_request_serialization_string_input() {
+    let req = EmbeddingRequest {
+        model: "test-emb".to_string(),
+        input: json!("hello"),
+        encoding_format: Some("float".to_string()),
+        user: Some("user-1".to_string()),
+        dimensions: Some(128),
+        rid: Some("rid-123".to_string()),
+        log_metrics: None,
+    };
+
+    let serialized = to_string(&req).unwrap();
+    let deserialized: EmbeddingRequest = from_str(&serialized).unwrap();
+
+    assert_eq!(deserialized.model, req.model);
+    assert_eq!(deserialized.input, req.input);
+    assert_eq!(deserialized.encoding_format, req.encoding_format);
+    assert_eq!(deserialized.user, req.user);
+    assert_eq!(deserialized.dimensions, req.dimensions);
+    assert_eq!(deserialized.rid, req.rid);
+}
+
+#[test]
+fn test_embedding_request_serialization_array_input() {
+    let req = EmbeddingRequest {
+        model: "test-emb".to_string(),
+        input: json!(["a", "b", "c"]),
+        encoding_format: None,
+        user: None,
+        dimensions: None,
+        rid: None,
+        log_metrics: None,
+    };
+
+    let serialized = to_string(&req).unwrap();
+    let de: EmbeddingRequest = from_str(&serialized).unwrap();
+    assert_eq!(de.model, req.model);
+    assert_eq!(de.input, req.input);
+}
+
+#[test]
+fn test_embedding_generation_request_trait_string() {
+    let req = EmbeddingRequest {
+        model: "emb-model".to_string(),
+        input: json!("hello"),
+        encoding_format: None,
+        user: None,
+        dimensions: None,
+        rid: None,
+        log_metrics: None,
+    };
+    assert!(!req.is_stream());
+    assert_eq!(req.get_model(), Some("emb-model"));
+    assert_eq!(req.extract_text_for_routing(), "hello");
+}
+
+#[test]
+fn test_embedding_generation_request_trait_array() {
+    let req = EmbeddingRequest {
+        model: "emb-model".to_string(),
+        input: json!(["hello", "world"]),
+        encoding_format: None,
+        user: None,
+        dimensions: None,
+        rid: None,
+        log_metrics: None,
+    };
+    assert_eq!(req.extract_text_for_routing(), "hello world");
+}
+
+#[test]
+fn test_embedding_generation_request_trait_non_text() {
+    let req = EmbeddingRequest {
+        model: "emb-model".to_string(),
+        input: json!({"tokens": [1, 2, 3]}),
+        encoding_format: None,
+        user: None,
+        dimensions: None,
+        rid: None,
+        log_metrics: None,
+    };
+    assert_eq!(req.extract_text_for_routing(), "");
+}
+
+#[test]
+fn test_embedding_generation_request_trait_mixed_array_ignores_nested() {
+    let req = EmbeddingRequest {
+        model: "emb-model".to_string(),
+        input: json!(["a", ["b", "c"], 123, {"k": "v"}]),
+        encoding_format: None,
+        user: None,
+        dimensions: None,
+        rid: None,
+        log_metrics: None,
+    };
+    // Only top-level string elements are extracted
+    assert_eq!(req.extract_text_for_routing(), "a");
+}

third_party/sglang/sgl-model-gateway/tests/spec/mod.rs

@@ -0,0 +1,9 @@
+// Protocol specification tests
+// These tests were originally in src/protocols/spec.rs and have been moved here
+// to reduce the size of that file and improve test organization.
+
+mod chat_completion;
+mod chat_message;
+mod embedding;
+mod rerank;
+mod responses;

third_party/sglang/sgl-model-gateway/tests/spec/rerank.rs

@@ -0,0 +1,569 @@
+use std::collections::HashMap;
+
+use serde_json::{from_str, to_string, Number, Value};
+use smg::protocols::{
+    common::{GenerationRequest, StringOrArray, UsageInfo},
+    rerank::{RerankRequest, RerankResponse, RerankResult, V1RerankReqInput},
+};
+use validator::Validate;
+
+#[test]
+fn test_rerank_request_serialization() {
+    let request = RerankRequest {
+        query: "test query".to_string(),
+        documents: vec!["doc1".to_string(), "doc2".to_string()],
+        model: "test-model".to_string(),
+        top_k: Some(5),
+        return_documents: true,
+        rid: Some(StringOrArray::String("req-123".to_string())),
+        user: Some("user-456".to_string()),
+    };
+
+    let serialized = to_string(&request).unwrap();
+    let deserialized: RerankRequest = from_str(&serialized).unwrap();
+
+    assert_eq!(deserialized.query, request.query);
+    assert_eq!(deserialized.documents, request.documents);
+    assert_eq!(deserialized.model, request.model);
+    assert_eq!(deserialized.top_k, request.top_k);
+    assert_eq!(deserialized.return_documents, request.return_documents);
+    assert_eq!(deserialized.rid, request.rid);
+    assert_eq!(deserialized.user, request.user);
+}
+
+#[test]
+fn test_rerank_request_deserialization_with_defaults() {
+    let json = r#"{
+        "query": "test query",
+        "documents": ["doc1", "doc2"]
+    }"#;
+
+    let request: RerankRequest = from_str(json).unwrap();
+
+    assert_eq!(request.query, "test query");
+    assert_eq!(request.documents, vec!["doc1", "doc2"]);
+    assert_eq!(request.model, "unknown");
+    assert_eq!(request.top_k, None);
+    assert!(request.return_documents);
+    assert_eq!(request.rid, None);
+    assert_eq!(request.user, None);
+}
+
+#[test]
+fn test_rerank_request_validation_success() {
+    let request = RerankRequest {
+        query: "valid query".to_string(),
+        documents: vec!["doc1".to_string(), "doc2".to_string()],
+        model: "test-model".to_string(),
+        top_k: Some(2),
+        return_documents: true,
+        rid: None,
+        user: None,
+    };
+
+    assert!(request.validate().is_ok());
+}
+
+#[test]
+fn test_rerank_request_validation_empty_query() {
+    let request = RerankRequest {
+        query: "".to_string(),
+        documents: vec!["doc1".to_string()],
+        model: "test-model".to_string(),
+        top_k: None,
+        return_documents: true,
+        rid: None,
+        user: None,
+    };
+
+    let result = request.validate();
+    assert!(result.is_err(), "Should reject empty query");
+}
+
+#[test]
+fn test_rerank_request_validation_whitespace_query() {
+    let request = RerankRequest {
+        query: "   ".to_string(),
+        documents: vec!["doc1".to_string()],
+        model: "test-model".to_string(),
+        top_k: None,
+        return_documents: true,
+        rid: None,
+        user: None,
+    };
+
+    let result = request.validate();
+    assert!(result.is_err(), "Should reject whitespace-only query");
+}
+
+#[test]
+fn test_rerank_request_validation_empty_documents() {
+    let request = RerankRequest {
+        query: "test query".to_string(),
+        documents: vec![],
+        model: "test-model".to_string(),
+        top_k: None,
+        return_documents: true,
+        rid: None,
+        user: None,
+    };
+
+    let result = request.validate();
+    assert!(result.is_err(), "Should reject empty documents list");
+}
+
+#[test]
+fn test_rerank_request_validation_top_k_zero() {
+    let request = RerankRequest {
+        query: "test query".to_string(),
+        documents: vec!["doc1".to_string(), "doc2".to_string()],
+        model: "test-model".to_string(),
+        top_k: Some(0),
+        return_documents: true,
+        rid: None,
+        user: None,
+    };
+
+    let result = request.validate();
+    assert!(result.is_err(), "Should reject top_k of zero");
+}
+
+#[test]
+fn test_rerank_request_validation_top_k_greater_than_docs() {
+    let request = RerankRequest {
+        query: "test query".to_string(),
+        documents: vec!["doc1".to_string(), "doc2".to_string()],
+        model: "test-model".to_string(),
+        top_k: Some(5),
+        return_documents: true,
+        rid: None,
+        user: None,
+    };
+
+    // This should pass but log a warning
+    assert!(request.validate().is_ok());
+}
+
+#[test]
+fn test_rerank_request_effective_top_k() {
+    let request = RerankRequest {
+        query: "test query".to_string(),
+        documents: vec!["doc1".to_string(), "doc2".to_string(), "doc3".to_string()],
+        model: "test-model".to_string(),
+        top_k: Some(2),
+        return_documents: true,
+        rid: None,
+        user: None,
+    };
+
+    assert_eq!(request.effective_top_k(), 2);
+}
+
+#[test]
+fn test_rerank_request_effective_top_k_none() {
+    let request = RerankRequest {
+        query: "test query".to_string(),
+        documents: vec!["doc1".to_string(), "doc2".to_string(), "doc3".to_string()],
+        model: "test-model".to_string(),
+        top_k: None,
+        return_documents: true,
+        rid: None,
+        user: None,
+    };
+
+    assert_eq!(request.effective_top_k(), 3);
+}
+
+#[test]
+fn test_rerank_response_creation() {
+    let results = vec![
+        RerankResult {
+            score: 0.8,
+            document: Some("doc1".to_string()),
+            index: 0,
+            meta_info: None,
+        },
+        RerankResult {
+            score: 0.6,
+            document: Some("doc2".to_string()),
+            index: 1,
+            meta_info: None,
+        },
+    ];
+
+    let response = RerankResponse::new(
+        results.clone(),
+        "test-model".to_string(),
+        Some(StringOrArray::String("req-123".to_string())),
+    );
+
+    assert_eq!(response.results.len(), 2);
+    assert_eq!(response.model, "test-model");
+    assert_eq!(
+        response.id,
+        Some(StringOrArray::String("req-123".to_string()))
+    );
+    assert_eq!(response.object, "rerank");
+    assert!(response.created > 0);
+}
+
+#[test]
+fn test_rerank_response_serialization() {
+    let results = vec![RerankResult {
+        score: 0.8,
+        document: Some("doc1".to_string()),
+        index: 0,
+        meta_info: None,
+    }];
+
+    let response = RerankResponse::new(
+        results,
+        "test-model".to_string(),
+        Some(StringOrArray::String("req-123".to_string())),
+    );
+
+    let serialized = to_string(&response).unwrap();
+    let deserialized: RerankResponse = from_str(&serialized).unwrap();
+
+    assert_eq!(deserialized.results.len(), response.results.len());
+    assert_eq!(deserialized.model, response.model);
+    assert_eq!(deserialized.id, response.id);
+    assert_eq!(deserialized.object, response.object);
+}
+
+#[test]
+fn test_rerank_response_apply_top_k() {
+    let results = vec![
+        RerankResult {
+            score: 0.8,
+            document: Some("doc1".to_string()),
+            index: 0,
+            meta_info: None,
+        },
+        RerankResult {
+            score: 0.6,
+            document: Some("doc2".to_string()),
+            index: 1,
+            meta_info: None,
+        },
+        RerankResult {
+            score: 0.4,
+            document: Some("doc3".to_string()),
+            index: 2,
+            meta_info: None,
+        },
+    ];
+
+    let mut response = RerankResponse::new(
+        results,
+        "test-model".to_string(),
+        Some(StringOrArray::String("req-123".to_string())),
+    );
+
+    response.apply_top_k(2);
+
+    assert_eq!(response.results.len(), 2);
+    assert_eq!(response.results[0].score, 0.8);
+    assert_eq!(response.results[1].score, 0.6);
+}
+
+#[test]
+fn test_rerank_response_apply_top_k_larger_than_results() {
+    let results = vec![RerankResult {
+        score: 0.8,
+        document: Some("doc1".to_string()),
+        index: 0,
+        meta_info: None,
+    }];
+
+    let mut response = RerankResponse::new(
+        results,
+        "test-model".to_string(),
+        Some(StringOrArray::String("req-123".to_string())),
+    );
+
+    response.apply_top_k(5);
+
+    assert_eq!(response.results.len(), 1);
+}
+
+#[test]
+fn test_rerank_response_drop_documents() {
+    let results = vec![RerankResult {
+        score: 0.8,
+        document: Some("doc1".to_string()),
+        index: 0,
+        meta_info: None,
+    }];
+    let mut response = RerankResponse::new(
+        results,
+        "test-model".to_string(),
+        Some(StringOrArray::String("req-123".to_string())),
+    );
+
+    response.drop_documents();
+
+    assert_eq!(response.results[0].document, None);
+}
+
+#[test]
+fn test_rerank_result_serialization() {
+    let result = RerankResult {
+        score: 0.85,
+        document: Some("test document".to_string()),
+        index: 42,
+        meta_info: Some(HashMap::from([
+            ("confidence".to_string(), Value::String("high".to_string())),
+            (
+                "processing_time".to_string(),
+                Value::Number(Number::from(150)),
+            ),
+        ])),
+    };
+
+    let serialized = to_string(&result).unwrap();
+    let deserialized: RerankResult = from_str(&serialized).unwrap();
+
+    assert_eq!(deserialized.score, result.score);
+    assert_eq!(deserialized.document, result.document);
+    assert_eq!(deserialized.index, result.index);
+    assert_eq!(deserialized.meta_info, result.meta_info);
+}
+
+#[test]
+fn test_rerank_result_serialization_without_document() {
+    let result = RerankResult {
+        score: 0.85,
+        document: None,
+        index: 42,
+        meta_info: None,
+    };
+
+    let serialized = to_string(&result).unwrap();
+    let deserialized: RerankResult = from_str(&serialized).unwrap();
+
+    assert_eq!(deserialized.score, result.score);
+    assert_eq!(deserialized.document, result.document);
+    assert_eq!(deserialized.index, result.index);
+    assert_eq!(deserialized.meta_info, result.meta_info);
+}
+
+#[test]
+fn test_v1_rerank_req_input_serialization() {
+    let v1_input = V1RerankReqInput {
+        query: "test query".to_string(),
+        documents: vec!["doc1".to_string(), "doc2".to_string()],
+    };
+
+    let serialized = to_string(&v1_input).unwrap();
+    let deserialized: V1RerankReqInput = from_str(&serialized).unwrap();
+
+    assert_eq!(deserialized.query, v1_input.query);
+    assert_eq!(deserialized.documents, v1_input.documents);
+}
+
+#[test]
+fn test_v1_to_rerank_request_conversion() {
+    let v1_input = V1RerankReqInput {
+        query: "test query".to_string(),
+        documents: vec!["doc1".to_string(), "doc2".to_string()],
+    };
+
+    let request: RerankRequest = v1_input.into();
+
+    assert_eq!(request.query, "test query");
+    assert_eq!(request.documents, vec!["doc1", "doc2"]);
+    assert_eq!(request.model, "unknown");
+    assert_eq!(request.top_k, None);
+    assert!(request.return_documents);
+    assert_eq!(request.rid, None);
+    assert_eq!(request.user, None);
+}
+
+#[test]
+fn test_rerank_request_generation_request_trait() {
+    let request = RerankRequest {
+        query: "test query".to_string(),
+        documents: vec!["doc1".to_string()],
+        model: "test-model".to_string(),
+        top_k: None,
+        return_documents: true,
+        rid: None,
+        user: None,
+    };
+
+    assert_eq!(request.get_model(), Some("test-model"));
+    assert!(!request.is_stream());
+    assert_eq!(request.extract_text_for_routing(), "test query");
+}
+
+#[test]
+fn test_rerank_request_very_long_query() {
+    let long_query = "a".repeat(100000);
+    let request = RerankRequest {
+        query: long_query,
+        documents: vec!["doc1".to_string()],
+        model: "test-model".to_string(),
+        top_k: None,
+        return_documents: true,
+        rid: None,
+        user: None,
+    };
+
+    assert!(request.validate().is_ok());
+}
+
+#[test]
+fn test_rerank_request_many_documents() {
+    let documents: Vec<String> = (0..1000).map(|i| format!("doc{}", i)).collect();
+    let request = RerankRequest {
+        query: "test query".to_string(),
+        documents,
+        model: "test-model".to_string(),
+        top_k: Some(100),
+        return_documents: true,
+        rid: None,
+        user: None,
+    };
+
+    assert!(request.validate().is_ok());
+    assert_eq!(request.effective_top_k(), 100);
+}
+
+#[test]
+fn test_rerank_request_special_characters() {
+    let request = RerankRequest {
+        query: "query with émojis 🚀 and unicode: 测试".to_string(),
+        documents: vec![
+            "doc with émojis 🎉".to_string(),
+            "doc with unicode: 测试".to_string(),
+        ],
+        model: "test-model".to_string(),
+        top_k: None,
+        return_documents: true,
+        rid: Some(StringOrArray::String("req-🚀-123".to_string())),
+        user: Some("user-🎉-456".to_string()),
+    };
+
+    assert!(request.validate().is_ok());
+}
+
+#[test]
+fn test_rerank_request_rid_array() {
+    let request = RerankRequest {
+        query: "test query".to_string(),
+        documents: vec!["doc1".to_string()],
+        model: "test-model".to_string(),
+        top_k: None,
+        return_documents: true,
+        rid: Some(StringOrArray::Array(vec![
+            "req1".to_string(),
+            "req2".to_string(),
+        ])),
+        user: None,
+    };
+
+    assert!(request.validate().is_ok());
+}
+
+#[test]
+fn test_rerank_response_with_usage_info() {
+    let results = vec![RerankResult {
+        score: 0.8,
+        document: Some("doc1".to_string()),
+        index: 0,
+        meta_info: None,
+    }];
+
+    let mut response = RerankResponse::new(
+        results,
+        "test-model".to_string(),
+        Some(StringOrArray::String("req-123".to_string())),
+    );
+
+    response.usage = Some(UsageInfo {
+        prompt_tokens: 100,
+        completion_tokens: 50,
+        total_tokens: 150,
+        reasoning_tokens: None,
+        prompt_tokens_details: None,
+    });
+
+    let serialized = to_string(&response).unwrap();
+    let deserialized: RerankResponse = from_str(&serialized).unwrap();
+
+    assert!(deserialized.usage.is_some());
+    let usage = deserialized.usage.unwrap();
+    assert_eq!(usage.prompt_tokens, 100);
+    assert_eq!(usage.completion_tokens, 50);
+    assert_eq!(usage.total_tokens, 150);
+}
+
+#[test]
+fn test_full_rerank_workflow() {
+    // Create request
+    let request = RerankRequest {
+        query: "machine learning".to_string(),
+        documents: vec![
+            "Introduction to machine learning algorithms".to_string(),
+            "Deep learning for computer vision".to_string(),
+            "Natural language processing basics".to_string(),
+            "Statistics and probability theory".to_string(),
+        ],
+        model: "rerank-model".to_string(),
+        top_k: Some(2),
+        return_documents: true,
+        rid: Some(StringOrArray::String("req-123".to_string())),
+        user: Some("user-456".to_string()),
+    };
+
+    // Validate request
+    assert!(request.validate().is_ok());
+
+    // Simulate reranking results (in real scenario, this would come from the model)
+    let results = vec![
+        RerankResult {
+            score: 0.95,
+            document: Some("Introduction to machine learning algorithms".to_string()),
+            index: 0,
+            meta_info: None,
+        },
+        RerankResult {
+            score: 0.87,
+            document: Some("Deep learning for computer vision".to_string()),
+            index: 1,
+            meta_info: None,
+        },
+        RerankResult {
+            score: 0.72,
+            document: Some("Natural language processing basics".to_string()),
+            index: 2,
+            meta_info: None,
+        },
+        RerankResult {
+            score: 0.45,
+            document: Some("Statistics and probability theory".to_string()),
+            index: 3,
+            meta_info: None,
+        },
+    ];
+
+    // Create response
+    let mut response = RerankResponse::new(results, request.model.clone(), request.rid.clone());
+
+    // Apply top_k
+    response.apply_top_k(request.effective_top_k());
+
+    assert_eq!(response.results.len(), 2);
+    assert_eq!(response.results[0].score, 0.95);
+    assert_eq!(response.results[0].index, 0);
+    assert_eq!(response.results[1].score, 0.87);
+    assert_eq!(response.results[1].index, 1);
+    assert_eq!(response.model, "rerank-model");
+
+    // Serialize and deserialize
+    let serialized = to_string(&response).unwrap();
+    let deserialized: RerankResponse = from_str(&serialized).unwrap();
+    assert_eq!(deserialized.results.len(), 2);
+    assert_eq!(deserialized.model, response.model);
+}

third_party/sglang/sgl-model-gateway/tests/spec_test.rs

@@ -0,0 +1,2 @@
+// Protocol specification tests
+mod spec;

third_party/sglang/sgl-model-gateway/tests/wasm_test.rs

@@ -0,0 +1,804 @@
+//! WASM Module Integration Tests
+//!
+//! This test suite validates the complete WASM module management functionality:
+//! - API endpoints (add, remove, list)
+//! - Workflow integration
+//! - Module execution
+//! - Error handling
+
+mod common;
+
+use std::{sync::Arc, time::Duration};
+
+use axum::{
+    body::{to_bytes, Body},
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use data_connector::{
+    MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage,
+};
+use smg::{
+    app_context::AppContext,
+    config::RouterConfig,
+    core::{LoadMonitor, WorkerRegistry},
+    policies::PolicyRegistry,
+    routers::RouterFactory,
+    server::{build_app, AppState},
+    tokenizer::TokenizerRegistry,
+    wasm::{
+        module::{
+            WasmModuleAddRequest, WasmModuleAddResponse, WasmModuleAttachPoint,
+            WasmModuleDescriptor, WasmModuleListResponse, WasmModuleType,
+        },
+        module_manager::WasmModuleManager,
+    },
+};
+use tempfile::TempDir;
+use tokio::fs;
+use tower::ServiceExt;
+use uuid::Uuid;
+
+/// Create a test AppContext with WASM manager initialized
+async fn create_test_context_with_wasm() -> Arc<AppContext> {
+    let config = RouterConfig::default();
+
+    // Initialize WASM manager first
+    let wasm_manager = Arc::new(
+        WasmModuleManager::with_default_config().expect("Failed to create WASM module manager"),
+    );
+
+    // Create AppContext with wasm_manager from the start
+    let client = reqwest::Client::new();
+
+    let tokenizer_registry = Arc::new(TokenizerRegistry::new());
+    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 engines, and mcp manager
+    use std::sync::OnceLock;
+    let worker_job_queue = Arc::new(OnceLock::new());
+    let workflow_engines = Arc::new(OnceLock::new());
+    let mcp_manager_lock = Arc::new(OnceLock::new());
+
+    let app_context = Arc::new(
+        AppContext::builder()
+            .router_config(config.clone())
+            .client(client)
+            .rate_limiter(None)
+            .tokenizer_registry(tokenizer_registry)
+            .reasoning_parser_factory(None)
+            .tool_parser_factory(None)
+            .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_engines(workflow_engines)
+            .mcp_manager(mcp_manager_lock)
+            .wasm_manager(Some(wasm_manager))
+            .build()
+            .expect("Failed to build AppContext with WASM manager"),
+    );
+
+    // Initialize JobQueue after AppContext is created
+    let weak_context = Arc::downgrade(&app_context);
+    let job_queue = smg::core::JobQueue::new(smg::core::JobQueueConfig::default(), weak_context);
+    app_context
+        .worker_job_queue
+        .set(job_queue)
+        .expect("JobQueue should only be initialized once");
+
+    // Initialize WorkflowEngines
+    use smg::core::steps::WorkflowEngines;
+    let engines = WorkflowEngines::new(&config);
+    app_context
+        .workflow_engines
+        .set(engines)
+        .expect("WorkflowEngines should only be initialized once");
+
+    // Initialize MCP manager with empty config
+    use smg_mcp::{McpConfig, McpManager};
+    let empty_config = McpConfig {
+        servers: vec![],
+        pool: Default::default(),
+        proxy: None,
+        warmup: vec![],
+        inventory: Default::default(),
+    };
+    let mcp_manager = McpManager::with_defaults(empty_config)
+        .await
+        .expect("Failed to create MCP manager");
+    app_context
+        .mcp_manager
+        .set(Arc::new(mcp_manager))
+        .ok()
+        .expect("McpManager should only be initialized once");
+
+    app_context
+}
+
+/// Create a test WASM component file
+/// Dynamically generates a valid WASM component programmatically without external tools
+/// This ensures tests work in new environments without requiring pre-built files or external tools
+async fn create_test_wasm_component(temp_dir: &TempDir) -> String {
+    use wasm_encoder::{Component, Module};
+
+    // Create a minimal valid WASM module first
+    // A minimal module needs at least a type section
+    let mut module = Module::new();
+
+    // Add an empty type section (0 types) - this is valid
+    let type_section = wasm_encoder::TypeSection::new();
+    module.section(&type_section);
+    let mut component = Component::new();
+    component.section(&wasm_encoder::ModuleSection(&module));
+    let component_bytes = component.as_slice().to_vec();
+    let component_path = temp_dir.path().join("test_module.component.wasm");
+    fs::write(&component_path, component_bytes)
+        .await
+        .expect("Failed to write WASM component file");
+
+    // Return absolute path
+    component_path
+        .canonicalize()
+        .expect("Failed to canonicalize path")
+        .to_str()
+        .unwrap()
+        .to_string()
+}
+
+/// Create a test app with WASM support
+async fn create_test_app_with_wasm() -> (axum::Router, Arc<AppContext>, TempDir) {
+    let temp_dir = TempDir::new().expect("Failed to create temp directory");
+    let app_context = create_test_context_with_wasm().await;
+
+    // Create a dummy router (we only need the app for WASM endpoints)
+    let router = RouterFactory::create_router(&app_context)
+        .await
+        .expect("Failed to create router");
+    let router = Arc::from(router);
+
+    let app_state = Arc::new(AppState {
+        router,
+        context: app_context.clone(),
+        concurrency_queue_tx: None,
+        router_manager: None,
+        mesh_handler: None,
+        mesh_sync_manager: None,
+    });
+
+    let request_id_headers = vec!["x-request-id".to_string(), "x-correlation-id".to_string()];
+
+    let app = build_app(
+        app_state,
+        smg::middleware::AuthConfig { api_key: None },
+        None, // No control plane auth for tests
+        256 * 1024 * 1024,
+        request_id_headers,
+        vec![], // cors_allowed_origins
+    );
+
+    (app, app_context, temp_dir)
+}
+
+// ============================================================================
+// API Endpoint Tests
+// ============================================================================
+
+#[tokio::test]
+async fn test_wasm_api_add_module() {
+    let (app, app_context, temp_dir) = create_test_app_with_wasm().await;
+    let wasm_file_path = create_test_wasm_component(&temp_dir).await;
+
+    let add_request = WasmModuleAddRequest {
+        modules: vec![WasmModuleDescriptor {
+            name: "test_module".to_string(),
+            file_path: wasm_file_path.clone(),
+            module_type: WasmModuleType::Middleware,
+            attach_points: vec![WasmModuleAttachPoint::Middleware(
+                smg::wasm::module::MiddlewareAttachPoint::OnRequest,
+            )],
+            add_result: None,
+        }],
+    };
+
+    let response = app
+        .oneshot(
+            Request::builder()
+                .method("POST")
+                .uri("/wasm")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&add_request).unwrap()))
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    let status = response.status();
+    let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
+    let response_json: WasmModuleAddResponse = serde_json::from_slice(&body).unwrap();
+
+    assert_eq!(response_json.modules.len(), 1);
+    let module_result = &response_json.modules[0].add_result;
+
+    // Print error for debugging
+    if let Some(smg::wasm::module::WasmModuleAddResult::Error(err)) = module_result {
+        eprintln!("Module registration failed: {}", err);
+    }
+
+    // If status is not OK, check the error message
+    if status != StatusCode::OK {
+        eprintln!("Response status: {:?}", status);
+        eprintln!("Response body: {}", String::from_utf8_lossy(&body));
+        panic!(
+            "Expected OK status but got {:?}. Error: {:?}",
+            status, module_result
+        );
+    }
+
+    assert!(module_result.is_some());
+
+    // Verify module is registered in wasm_manager
+    if let Some(wasm_manager) = app_context.wasm_manager.as_ref() {
+        let modules = wasm_manager.get_modules().expect("Failed to get modules");
+        assert!(!modules.is_empty(), "Module should be registered");
+
+        if let Some(smg::wasm::module::WasmModuleAddResult::Success(uuid)) = module_result {
+            let module = wasm_manager
+                .get_module(*uuid)
+                .expect("Failed to get module");
+            assert!(module.is_some(), "Module should exist in manager");
+        }
+    }
+}
+
+#[tokio::test]
+async fn test_wasm_api_add_module_invalid_file() {
+    let (app, _app_context, _temp_dir) = create_test_app_with_wasm().await;
+
+    let add_request = WasmModuleAddRequest {
+        modules: vec![WasmModuleDescriptor {
+            name: "test_module".to_string(),
+            file_path: "/nonexistent/path/to/module.component.wasm".to_string(),
+            module_type: WasmModuleType::Middleware,
+            attach_points: vec![WasmModuleAttachPoint::Middleware(
+                smg::wasm::module::MiddlewareAttachPoint::OnRequest,
+            )],
+            add_result: None,
+        }],
+    };
+
+    let response = app
+        .oneshot(
+            Request::builder()
+                .method("POST")
+                .uri("/wasm")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&add_request).unwrap()))
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    // Should return error status
+    assert!(
+        response.status() == StatusCode::BAD_REQUEST
+            || response.status() == StatusCode::INTERNAL_SERVER_ERROR
+    );
+
+    let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
+    let response_json: WasmModuleAddResponse = serde_json::from_slice(&body).unwrap();
+
+    assert_eq!(response_json.modules.len(), 1);
+    let module_result = &response_json.modules[0].add_result;
+    assert!(module_result.is_some());
+
+    // Verify it's an error result
+    if let Some(smg::wasm::module::WasmModuleAddResult::Error(_)) = module_result {
+        // Expected error
+    } else {
+        panic!("Expected error result for invalid file path");
+    }
+}
+
+#[tokio::test]
+async fn test_wasm_api_add_module_invalid_wasm() {
+    let (app, _app_context, temp_dir) = create_test_app_with_wasm().await;
+
+    // Create an invalid WASM file (just random bytes)
+    let invalid_wasm_path = temp_dir.path().join("invalid.component.wasm");
+    fs::write(&invalid_wasm_path, b"not a valid wasm file")
+        .await
+        .expect("Failed to write invalid WASM file");
+
+    let add_request = WasmModuleAddRequest {
+        modules: vec![WasmModuleDescriptor {
+            name: "invalid_module".to_string(),
+            file_path: invalid_wasm_path.to_str().unwrap().to_string(),
+            module_type: WasmModuleType::Middleware,
+            attach_points: vec![WasmModuleAttachPoint::Middleware(
+                smg::wasm::module::MiddlewareAttachPoint::OnRequest,
+            )],
+            add_result: None,
+        }],
+    };
+
+    let response = app
+        .oneshot(
+            Request::builder()
+                .method("POST")
+                .uri("/wasm")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&add_request).unwrap()))
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    // Should return error status
+    assert!(
+        response.status() == StatusCode::BAD_REQUEST
+            || response.status() == StatusCode::INTERNAL_SERVER_ERROR
+    );
+
+    let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
+    let response_json: WasmModuleAddResponse = serde_json::from_slice(&body).unwrap();
+
+    assert_eq!(response_json.modules.len(), 1);
+    let module_result = &response_json.modules[0].add_result;
+    assert!(module_result.is_some());
+
+    // Verify it's an error result
+    if let Some(smg::wasm::module::WasmModuleAddResult::Error(_)) = module_result {
+        // Expected error
+    } else {
+        panic!("Expected error result for invalid WASM file");
+    }
+}
+
+#[tokio::test]
+async fn test_wasm_api_list_modules() {
+    let (app, _app_context, temp_dir) = create_test_app_with_wasm().await;
+    let wasm_file_path = create_test_wasm_component(&temp_dir).await;
+
+    // First, add a module
+    let add_request = WasmModuleAddRequest {
+        modules: vec![WasmModuleDescriptor {
+            name: "test_module_list".to_string(),
+            file_path: wasm_file_path.clone(),
+            module_type: WasmModuleType::Middleware,
+            attach_points: vec![WasmModuleAttachPoint::Middleware(
+                smg::wasm::module::MiddlewareAttachPoint::OnRequest,
+            )],
+            add_result: None,
+        }],
+    };
+
+    let add_response = app
+        .clone()
+        .oneshot(
+            Request::builder()
+                .method("POST")
+                .uri("/wasm")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&add_request).unwrap()))
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    assert_eq!(add_response.status(), StatusCode::OK);
+
+    // Wait a bit for the job to complete
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    // Now list modules
+    let list_response = app
+        .oneshot(
+            Request::builder()
+                .method("GET")
+                .uri("/wasm")
+                .body(Body::empty())
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    assert_eq!(list_response.status(), StatusCode::OK);
+
+    let body = to_bytes(list_response.into_body(), usize::MAX)
+        .await
+        .unwrap();
+    let response_json: WasmModuleListResponse = serde_json::from_slice(&body).unwrap();
+
+    assert!(
+        !response_json.modules.is_empty(),
+        "Should have at least one module"
+    );
+    assert!(response_json
+        .modules
+        .iter()
+        .any(|m| m.module_meta.name == "test_module_list"));
+
+    // Verify metrics are present (total_executions is u64, so always >= 0)
+    let _ = response_json.metrics.total_executions;
+}
+
+#[tokio::test]
+async fn test_wasm_api_remove_module() {
+    let (app, app_context, temp_dir) = create_test_app_with_wasm().await;
+    let wasm_file_path = create_test_wasm_component(&temp_dir).await;
+
+    // First, add a module
+    let add_request = WasmModuleAddRequest {
+        modules: vec![WasmModuleDescriptor {
+            name: "test_module_remove".to_string(),
+            file_path: wasm_file_path.clone(),
+            module_type: WasmModuleType::Middleware,
+            attach_points: vec![WasmModuleAttachPoint::Middleware(
+                smg::wasm::module::MiddlewareAttachPoint::OnRequest,
+            )],
+            add_result: None,
+        }],
+    };
+
+    let add_response = app
+        .clone()
+        .oneshot(
+            Request::builder()
+                .method("POST")
+                .uri("/wasm")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&add_request).unwrap()))
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    assert_eq!(add_response.status(), StatusCode::OK);
+
+    let body = to_bytes(add_response.into_body(), usize::MAX)
+        .await
+        .unwrap();
+    let response_json: WasmModuleAddResponse = serde_json::from_slice(&body).unwrap();
+
+    // Wait for job to complete
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    // Get the module UUID
+    let module_uuid = if let Some(smg::wasm::module::WasmModuleAddResult::Success(uuid)) =
+        &response_json.modules[0].add_result
+    {
+        *uuid
+    } else {
+        // If we can't get UUID from response, try to find it from manager
+        if let Some(wasm_manager) = app_context.wasm_manager.as_ref() {
+            let modules = wasm_manager.get_modules().expect("Failed to get modules");
+            modules
+                .iter()
+                .find(|m| m.module_meta.name == "test_module_remove")
+                .map(|m| m.module_uuid)
+                .expect("Module should be registered")
+        } else {
+            panic!("WASM manager not available");
+        }
+    };
+
+    // Now remove the module
+    let remove_response = app
+        .oneshot(
+            Request::builder()
+                .method("DELETE")
+                .uri(format!("/wasm/{}", module_uuid))
+                .body(Body::empty())
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    let remove_status = remove_response.status();
+    if remove_status != StatusCode::OK {
+        let body = to_bytes(remove_response.into_body(), usize::MAX)
+            .await
+            .unwrap();
+        eprintln!(
+            "Remove module failed with status {:?}: {}",
+            remove_status,
+            String::from_utf8_lossy(&body)
+        );
+        panic!("Expected OK status but got {:?}", remove_status);
+    }
+
+    // Wait for removal to complete
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    // Verify module is removed
+    if let Some(wasm_manager) = app_context.wasm_manager.as_ref() {
+        let module = wasm_manager
+            .get_module(module_uuid)
+            .expect("Failed to get module");
+        assert!(module.is_none(), "Module should be removed");
+    }
+}
+
+#[tokio::test]
+async fn test_wasm_api_remove_module_not_found() {
+    let (app, _app_context, _temp_dir) = create_test_app_with_wasm().await;
+
+    // Try to remove a non-existent module
+    let fake_uuid = Uuid::new_v4();
+    let remove_response = app
+        .oneshot(
+            Request::builder()
+                .method("DELETE")
+                .uri(format!("/wasm/{}", fake_uuid))
+                .body(Body::empty())
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    // Should return error status
+    assert!(
+        remove_response.status() == StatusCode::BAD_REQUEST
+            || remove_response.status() == StatusCode::NOT_FOUND
+    );
+}
+
+// ============================================================================
+// WASM Functionality Tests
+// ============================================================================
+
+#[tokio::test]
+async fn test_wasm_module_duplicate_sha256() {
+    let (app, _app_context, temp_dir) = create_test_app_with_wasm().await;
+    let wasm_file_path = create_test_wasm_component(&temp_dir).await;
+
+    // Add first module
+    let add_request1 = WasmModuleAddRequest {
+        modules: vec![WasmModuleDescriptor {
+            name: "test_module_dup1".to_string(),
+            file_path: wasm_file_path.clone(),
+            module_type: WasmModuleType::Middleware,
+            attach_points: vec![WasmModuleAttachPoint::Middleware(
+                smg::wasm::module::MiddlewareAttachPoint::OnRequest,
+            )],
+            add_result: None,
+        }],
+    };
+
+    let response1 = app
+        .clone()
+        .oneshot(
+            Request::builder()
+                .method("POST")
+                .uri("/wasm")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&add_request1).unwrap()))
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    assert_eq!(response1.status(), StatusCode::OK);
+
+    // Wait for first job to complete
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    // Try to add the same file again (should fail due to duplicate SHA256)
+    let add_request2 = WasmModuleAddRequest {
+        modules: vec![WasmModuleDescriptor {
+            name: "test_module_dup2".to_string(),
+            file_path: wasm_file_path.clone(), // Same file
+            module_type: WasmModuleType::Middleware,
+            attach_points: vec![WasmModuleAttachPoint::Middleware(
+                smg::wasm::module::MiddlewareAttachPoint::OnRequest,
+            )],
+            add_result: None,
+        }],
+    };
+
+    let response2 = app
+        .oneshot(
+            Request::builder()
+                .method("POST")
+                .uri("/wasm")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&add_request2).unwrap()))
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    // Should return error status for duplicate
+    assert!(
+        response2.status() == StatusCode::BAD_REQUEST
+            || response2.status() == StatusCode::INTERNAL_SERVER_ERROR
+    );
+
+    let body = to_bytes(response2.into_body(), usize::MAX).await.unwrap();
+    let response_json: WasmModuleAddResponse = serde_json::from_slice(&body).unwrap();
+
+    assert_eq!(response_json.modules.len(), 1);
+    let module_result = &response_json.modules[0].add_result;
+    assert!(module_result.is_some());
+
+    // Verify it's an error result (duplicate)
+    if let Some(smg::wasm::module::WasmModuleAddResult::Error(err_msg)) = module_result {
+        assert!(
+            err_msg.contains("duplicate")
+                || err_msg.contains("Duplicate")
+                || err_msg.contains("SHA256")
+        );
+    } else {
+        panic!("Expected error result for duplicate SHA256");
+    }
+}
+
+#[tokio::test]
+async fn test_wasm_module_execution() {
+    let (_app, app_context, temp_dir) = create_test_app_with_wasm().await;
+    let wasm_file_path = create_test_wasm_component(&temp_dir).await;
+
+    // First, add a module using the workflow directly
+    let wasm_manager = app_context
+        .wasm_manager
+        .as_ref()
+        .expect("WASM manager should be initialized");
+
+    let engines = app_context
+        .workflow_engines
+        .get()
+        .expect("Workflow engines should be initialized");
+
+    // Create workflow context for registration
+    use smg::core::steps::{WasmModuleConfigRequest, WasmRegistrationWorkflowData};
+    use wfaas::WorkflowId;
+
+    let descriptor = WasmModuleDescriptor {
+        name: "test_execution_module".to_string(),
+        file_path: wasm_file_path.clone(),
+        module_type: WasmModuleType::Middleware,
+        attach_points: vec![WasmModuleAttachPoint::Middleware(
+            smg::wasm::module::MiddlewareAttachPoint::OnRequest,
+        )],
+        add_result: None,
+    };
+
+    let config_request = WasmModuleConfigRequest { descriptor };
+    let workflow_data = WasmRegistrationWorkflowData {
+        config: config_request,
+        wasm_bytes: None,
+        sha256_hash: None,
+        file_size_bytes: None,
+        module_uuid: None,
+        app_context: Some(app_context.clone()),
+    };
+
+    // Start workflow
+    let instance_id = engines
+        .wasm_registration
+        .start_workflow(WorkflowId::new("wasm_module_registration"), workflow_data)
+        .await
+        .expect("Failed to start workflow");
+
+    // Wait for workflow to complete
+    let timeout = Duration::from_secs(30);
+    let start = std::time::Instant::now();
+
+    let module_uuid = loop {
+        if start.elapsed() > timeout {
+            panic!("Workflow timeout");
+        }
+
+        let state = engines
+            .wasm_registration
+            .get_status(instance_id)
+            .await
+            .expect("Failed to get workflow status");
+
+        match state.status {
+            wfaas::WorkflowStatus::Completed => {
+                // Extract module UUID from typed workflow data
+                break state
+                    .context
+                    .data
+                    .module_uuid
+                    .expect("Module UUID should be in context");
+            }
+            wfaas::WorkflowStatus::Failed => {
+                panic!("Workflow failed: {:?}", state);
+            }
+            _ => {
+                tokio::time::sleep(Duration::from_millis(100)).await;
+            }
+        }
+    };
+
+    // Verify module is registered
+    let module = wasm_manager
+        .get_module(module_uuid)
+        .expect("Failed to get module");
+    assert!(module.is_some(), "Module should be registered");
+
+    // Get initial metrics
+    let (initial_total, initial_success, initial_failed, _, _) = wasm_manager.get_metrics();
+
+    // Execute the module
+    use smg::wasm::{
+        spec::smg::gateway::middleware_types,
+        types::{WasmComponentInput, WasmComponentOutput},
+    };
+
+    let request = middleware_types::Request {
+        method: "GET".to_string(),
+        path: "/test".to_string(),
+        query: "".to_string(),
+        headers: vec![],
+        body: vec![],
+        request_id: "test-request-id".to_string(),
+        now_epoch_ms: 1000,
+    };
+
+    let input = WasmComponentInput::MiddlewareRequest(request);
+    let attach_point =
+        WasmModuleAttachPoint::Middleware(smg::wasm::module::MiddlewareAttachPoint::OnRequest);
+
+    // Execute the module
+    let result = wasm_manager
+        .execute_module_interface(module_uuid, attach_point, input)
+        .await;
+
+    // Verify execution result
+    match result {
+        Ok(WasmComponentOutput::MiddlewareAction(action)) => {
+            // Verify action is valid (should be Continue, Reject, or Modify)
+            match action {
+                middleware_types::Action::Continue => {
+                    // Expected for a simple middleware
+                }
+                middleware_types::Action::Reject(_) => {
+                    // Also valid
+                }
+                middleware_types::Action::Modify(_) => {
+                    // Also valid
+                }
+            }
+        }
+        Err(e) => {
+            // Execution might fail if the WASM component is not properly built
+            // This is acceptable for testing - we're testing the execution path, not the component itself
+            eprintln!(
+                "Module execution failed (expected if component is not properly built): {:?}",
+                e
+            );
+        }
+    }
+
+    // Verify metrics were updated
+    let (final_total, final_success, final_failed, _, _) = wasm_manager.get_metrics();
+
+    // Metrics should have increased (either success or failed)
+    assert!(
+        final_total > initial_total
+            || final_failed > initial_failed
+            || final_success > initial_success,
+        "Metrics should be updated after execution"
+    );
+}