[model-gateway] change sgl-router to sgl-model-gateway (#14312)

sgl-model-gateway/tests/cache_aware_backward_compat_test.rs

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

sgl-model-gateway/tests/chat_template_format_detection.rs

@@ -0,0 +1,312 @@
+use sgl_model_gateway::{
+    protocols::chat::{ChatMessage, MessageContent},
+    tokenizer::chat_template::{
+        detect_chat_template_content_format, ChatTemplateContentFormat, ChatTemplateParams,
+        ChatTemplateProcessor,
+    },
+};
+
+#[test]
+fn test_detect_string_format_deepseek() {
+    // DeepSeek style template - expects string content
+    let template = r#"
+        {%- for message in messages %}
+        {%- if message['role'] == 'user' %}
+        User: {{ message['content'] }}
+        {%- elif message['role'] == 'assistant' %}
+        Assistant: {{ message['content'] }}
+        {%- endif %}
+        {%- endfor %}
+        "#;
+
+    assert_eq!(
+        detect_chat_template_content_format(template),
+        ChatTemplateContentFormat::String
+    );
+}
+
+#[test]
+fn test_detect_openai_format_llama4() {
+    // Llama4 style template - expects structured content
+    let template = r#"
+        {%- for message in messages %}
+        {%- if message['content'] is iterable %}
+        {%- for content in message['content'] %}
+        {%- if content['type'] == 'text' %}
+        {{ content['text'] }}
+        {%- elif content['type'] == 'image' %}
+        <image>
+        {%- endif %}
+        {%- endfor %}
+        {%- else %}
+        {{ message['content'] }}
+        {%- endif %}
+        {%- endfor %}
+        "#;
+
+    assert_eq!(
+        detect_chat_template_content_format(template),
+        ChatTemplateContentFormat::OpenAI
+    );
+}
+
+#[test]
+fn test_detect_openai_format_dot_notation() {
+    // Template using dot notation
+    let template = r#"
+        {%- for message in messages %}
+        {%- for part in message.content %}
+        {%- if part.type == 'text' %}
+        {{ part.text }}
+        {%- endif %}
+        {%- endfor %}
+        {%- endfor %}
+        "#;
+
+    assert_eq!(
+        detect_chat_template_content_format(template),
+        ChatTemplateContentFormat::OpenAI
+    );
+}
+
+#[test]
+fn test_detect_openai_format_variable_assignment() {
+    // Template that assigns content to variable then iterates
+    let template = r#"
+        {%- for message in messages %}
+        {%- set content = message['content'] %}
+        {%- if content is sequence %}
+        {%- for item in content %}
+        {{ item }}
+        {%- endfor %}
+        {%- endif %}
+        {%- endfor %}
+        "#;
+
+    assert_eq!(
+        detect_chat_template_content_format(template),
+        ChatTemplateContentFormat::OpenAI
+    );
+}
+
+#[test]
+fn test_detect_openai_format_glm4v_style() {
+    // GLM4V uses 'msg' instead of 'message'
+    let template = r#"
+        {%- for msg in messages %}
+        {%- for part in msg.content %}
+        {%- if part.type == 'text' %}{{ part.text }}{%- endif %}
+        {%- if part.type == 'image' %}<image>{%- endif %}
+        {%- endfor %}
+        {%- endfor %}
+        "#;
+
+    assert_eq!(
+        detect_chat_template_content_format(template),
+        ChatTemplateContentFormat::OpenAI
+    );
+}
+
+#[test]
+fn test_detect_openai_format_with_length_check() {
+    // Template that checks content length
+    let template = r#"
+        {%- for message in messages %}
+        {%- if message.content|length > 0 %}
+        {%- for item in message.content %}
+        {{ item.text }}
+        {%- endfor %}
+        {%- endif %}
+        {%- endfor %}
+        "#;
+
+    assert_eq!(
+        detect_chat_template_content_format(template),
+        ChatTemplateContentFormat::OpenAI
+    );
+}
+
+#[test]
+fn test_detect_openai_format_with_index_access() {
+    // Template that accesses content by index
+    let template = r#"
+        {%- for message in messages %}
+        {%- if message.content[0] %}
+        First item: {{ message.content[0].text }}
+        {%- endif %}
+        {%- endfor %}
+        "#;
+
+    assert_eq!(
+        detect_chat_template_content_format(template),
+        ChatTemplateContentFormat::OpenAI
+    );
+}
+
+#[test]
+fn test_invalid_template_defaults_to_string() {
+    let template = "Not a valid {% jinja template";
+
+    assert_eq!(
+        detect_chat_template_content_format(template),
+        ChatTemplateContentFormat::String
+    );
+}
+
+#[test]
+fn test_empty_template_defaults_to_string() {
+    assert_eq!(
+        detect_chat_template_content_format(""),
+        ChatTemplateContentFormat::String
+    );
+}
+
+#[test]
+fn test_simple_chat_template_unit_test() {
+    let template = r#"
+{%- for message in messages %}
+{{ message.role }}: {{ message.content }}
+{% endfor -%}
+{%- if add_generation_prompt %}
+assistant:
+{%- endif %}
+"#;
+
+    let processor = ChatTemplateProcessor::new(template.to_string());
+
+    let messages = [
+        ChatMessage::System {
+            content: MessageContent::Text("You are helpful".to_string()),
+            name: None,
+        },
+        ChatMessage::User {
+            content: MessageContent::Text("Hello".to_string()),
+            name: None,
+        },
+    ];
+
+    // Convert to JSON values like the router does
+    let message_values: Vec<serde_json::Value> = messages
+        .iter()
+        .map(|msg| serde_json::to_value(msg).unwrap())
+        .collect();
+
+    let params = ChatTemplateParams {
+        add_generation_prompt: true,
+        ..Default::default()
+    };
+    let result = processor
+        .apply_chat_template(&message_values, params)
+        .unwrap();
+    assert!(result.contains("system: You are helpful"));
+    assert!(result.contains("user: Hello"));
+    assert!(result.contains("assistant:"));
+}
+
+#[test]
+fn test_chat_template_with_tokens_unit_test() {
+    // Template that uses template kwargs for tokens (more realistic)
+    let template = r#"
+{%- if start_token -%}{{ start_token }}{%- endif -%}
+{%- for message in messages -%}
+{{ message.role }}: {{ message.content }}{%- if end_token -%}{{ end_token }}{%- endif -%}
+{% endfor -%}
+"#;
+
+    let processor = ChatTemplateProcessor::new(template.to_string());
+
+    let messages = [ChatMessage::User {
+        content: MessageContent::Text("Test".to_string()),
+        name: None,
+    }];
+
+    // Convert to JSON values like the router does
+    let message_values: Vec<serde_json::Value> = messages
+        .iter()
+        .map(|msg| serde_json::to_value(msg).unwrap())
+        .collect();
+
+    // Use template_kwargs to pass tokens
+    let mut template_kwargs = std::collections::HashMap::new();
+    template_kwargs.insert(
+        "start_token".to_string(),
+        serde_json::Value::String("<s>".to_string()),
+    );
+    template_kwargs.insert(
+        "end_token".to_string(),
+        serde_json::Value::String("</s>".to_string()),
+    );
+
+    let params = ChatTemplateParams {
+        template_kwargs: Some(&template_kwargs),
+        ..Default::default()
+    };
+
+    let result = processor
+        .apply_chat_template(&message_values, params)
+        .unwrap();
+    assert!(result.contains("<s>"));
+    assert!(result.contains("</s>"));
+}
+
+#[test]
+fn test_detect_openai_format_qwen3vl_macro_style() {
+    // Qwen3-VL style template using macros to handle multimodal content
+    // This tests the macro-based detection pattern
+    let template = r#"{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- else %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}
+                <|vision_start|><|image_pad|><|vision_end|>
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}
+                <|vision_start|><|video_pad|><|vision_end|>
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- endif %}
+        {%- endfor %}
+    {%- endif %}
+{%- endmacro %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, True) %}
+    {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}"#;
+
+    assert_eq!(
+        detect_chat_template_content_format(template),
+        ChatTemplateContentFormat::OpenAI
+    );
+}
+
+#[test]
+fn test_detect_openai_format_arbitrary_variable_names() {
+    // Test that detection works with any variable name, not just "message", "msg", "m"
+    // Uses "chat_msg" and "x" as loop variables
+    let template = r#"
+        {%- for chat_msg in messages %}
+        {%- for x in chat_msg.content %}
+        {%- if x.type == 'text' %}{{ x.text }}{%- endif %}
+        {%- if x.type == 'image' %}<image>{%- endif %}
+        {%- endfor %}
+        {%- endfor %}
+        "#;
+
+    assert_eq!(
+        detect_chat_template_content_format(template),
+        ChatTemplateContentFormat::OpenAI
+    );
+}

sgl-model-gateway/tests/chat_template_integration.rs

@@ -0,0 +1,414 @@
+use sgl_model_gateway::{
+    protocols::{
+        chat::{ChatMessage, MessageContent},
+        common::{ContentPart, ImageUrl},
+    },
+    tokenizer::chat_template::{
+        detect_chat_template_content_format, ChatTemplateContentFormat, ChatTemplateParams,
+        ChatTemplateProcessor,
+    },
+};
+
+#[test]
+fn test_simple_chat_template() {
+    let template = r#"
+{%- for message in messages %}
+<|{{ message.role }}|>{{ message.content }}<|end|>
+{% endfor -%}
+{%- if add_generation_prompt %}
+<|assistant|>
+{%- endif %}
+"#;
+
+    let processor = ChatTemplateProcessor::new(template.to_string());
+
+    let messages = [ChatMessage::User {
+        content: MessageContent::Text("Test".to_string()),
+        name: None,
+    }];
+
+    // Convert to JSON values like the router does
+    let message_values: Vec<serde_json::Value> = messages
+        .iter()
+        .map(|msg| serde_json::to_value(msg).unwrap())
+        .collect();
+
+    let params = ChatTemplateParams {
+        add_generation_prompt: true,
+        ..Default::default()
+    };
+    let result = processor
+        .apply_chat_template(&message_values, params)
+        .unwrap();
+    assert!(result.contains("<|user|>Test<|end|>"));
+    assert!(result.contains("<|assistant|>"));
+}
+
+#[test]
+fn test_chat_template_with_tokens() {
+    // Template that uses template kwargs for tokens
+    let template = r#"
+{%- if bos_token -%}{{ bos_token }}{%- endif -%}
+{%- for message in messages -%}
+{{ message.role }}: {{ message.content }}{%- if eos_token -%}{{ eos_token }}{%- endif -%}
+{% endfor -%}
+"#;
+
+    let processor = ChatTemplateProcessor::new(template.to_string());
+
+    let messages = [ChatMessage::User {
+        content: MessageContent::Text("Test".to_string()),
+        name: None,
+    }];
+
+    // Convert to JSON values like the router does
+    let message_values: Vec<serde_json::Value> = messages
+        .iter()
+        .map(|msg| serde_json::to_value(msg).unwrap())
+        .collect();
+
+    // Use template_kwargs to pass tokens
+    let mut template_kwargs = std::collections::HashMap::new();
+    template_kwargs.insert(
+        "bos_token".to_string(),
+        serde_json::Value::String("<s>".to_string()),
+    );
+    template_kwargs.insert(
+        "eos_token".to_string(),
+        serde_json::Value::String("</s>".to_string()),
+    );
+
+    let params = ChatTemplateParams {
+        template_kwargs: Some(&template_kwargs),
+        ..Default::default()
+    };
+
+    let result = processor
+        .apply_chat_template(&message_values, params)
+        .unwrap();
+    assert!(result.contains("<s>"));
+    assert!(result.contains("</s>"));
+}
+
+#[test]
+fn test_llama_style_template() {
+    let template = r#"
+{%- if messages[0]['role'] == 'system' -%}
+    {%- set system_message = messages[0]['content'] -%}
+    {%- set messages = messages[1:] -%}
+{%- else -%}
+    {%- set system_message = '' -%}
+{%- endif -%}
+
+{{- bos_token if bos_token else '<|begin_of_text|>' }}
+{%- if system_message %}
+{{- '<|start_header_id|>system<|end_header_id|>\n\n' + system_message + '<|eot_id|>' }}
+{%- endif %}
+
+{%- for message in messages %}
+    {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] + '<|eot_id|>' }}
+{%- endfor %}
+
+{%- if add_generation_prompt %}
+    {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}
+{%- endif %}
+"#;
+
+    let processor = ChatTemplateProcessor::new(template.to_string());
+
+    let messages = [
+        ChatMessage::System {
+            content: MessageContent::Text("You are a helpful assistant".to_string()),
+            name: None,
+        },
+        ChatMessage::User {
+            content: MessageContent::Text("What is 2+2?".to_string()),
+            name: None,
+        },
+    ];
+
+    // Convert to JSON values
+    let json_messages: Vec<serde_json::Value> = messages
+        .iter()
+        .map(|msg| serde_json::to_value(msg).unwrap())
+        .collect();
+
+    // Use template_kwargs to pass the token
+    let mut template_kwargs = std::collections::HashMap::new();
+    template_kwargs.insert(
+        "bos_token".to_string(),
+        serde_json::Value::String("<|begin_of_text|>".to_string()),
+    );
+
+    let params = ChatTemplateParams {
+        add_generation_prompt: true,
+        template_kwargs: Some(&template_kwargs),
+        ..Default::default()
+    };
+    let result = processor
+        .apply_chat_template(&json_messages, params)
+        .unwrap();
+
+    // Check that the result contains expected markers
+    assert!(result.contains("<|begin_of_text|>"));
+    assert!(result.contains("<|start_header_id|>system<|end_header_id|>"));
+    assert!(result.contains("You are a helpful assistant"));
+    assert!(result.contains("<|start_header_id|>user<|end_header_id|>"));
+    assert!(result.contains("What is 2+2?"));
+    assert!(result.contains("<|start_header_id|>assistant<|end_header_id|>"));
+}
+
+#[test]
+fn test_chatml_template() {
+    let template = r#"
+{%- for message in messages %}
+    {{- '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>\n' }}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}
+"#;
+
+    let processor = ChatTemplateProcessor::new(template.to_string());
+
+    let messages = [
+        ChatMessage::User {
+            content: MessageContent::Text("Hello".to_string()),
+            name: None,
+        },
+        ChatMessage::Assistant {
+            content: Some(MessageContent::Text("Hi there!".to_string())),
+            name: None,
+            tool_calls: None,
+            reasoning_content: None,
+        },
+        ChatMessage::User {
+            content: MessageContent::Text("How are you?".to_string()),
+            name: None,
+        },
+    ];
+
+    // Convert to JSON values
+    let json_messages: Vec<serde_json::Value> = messages
+        .iter()
+        .map(|msg| serde_json::to_value(msg).unwrap())
+        .collect();
+
+    let result = processor
+        .apply_chat_template(
+            &json_messages,
+            ChatTemplateParams {
+                add_generation_prompt: true,
+                ..Default::default()
+            },
+        )
+        .unwrap();
+
+    // Check ChatML format
+    assert!(result.contains("<|im_start|>user\nHello<|im_end|>"));
+    assert!(result.contains("<|im_start|>assistant\nHi there!<|im_end|>"));
+    assert!(result.contains("<|im_start|>user\nHow are you?<|im_end|>"));
+    assert!(result.ends_with("<|im_start|>assistant\n"));
+}
+
+#[test]
+fn test_template_without_generation_prompt() {
+    let template = r#"
+{%- for message in messages -%}
+{{ message.role }}: {{ message.content }}
+{% endfor -%}
+{%- if add_generation_prompt -%}
+assistant:
+{%- endif -%}
+"#;
+
+    let processor = ChatTemplateProcessor::new(template.to_string());
+
+    let messages = [ChatMessage::User {
+        content: MessageContent::Text("Test".to_string()),
+        name: None,
+    }];
+
+    // Convert to JSON values
+    let json_messages: Vec<serde_json::Value> = messages
+        .iter()
+        .map(|msg| serde_json::to_value(msg).unwrap())
+        .collect();
+
+    let result = processor
+        .apply_chat_template(&json_messages, ChatTemplateParams::default())
+        .unwrap();
+    assert_eq!(result.trim(), "user: Test");
+
+    let result_with_prompt = processor
+        .apply_chat_template(
+            &json_messages,
+            ChatTemplateParams {
+                add_generation_prompt: true,
+                ..Default::default()
+            },
+        )
+        .unwrap();
+    assert!(result_with_prompt.contains("assistant:"));
+}
+
+#[test]
+fn test_empty_messages_template() {
+    let template = r#"{% for msg in messages %}{{ msg.role }}: {{ msg.content }}\n{% endfor %}"#;
+
+    let processor = ChatTemplateProcessor::new(template.to_string());
+
+    let messages: Vec<serde_json::Value> = vec![];
+    let result = processor
+        .apply_chat_template(&messages, ChatTemplateParams::default())
+        .unwrap();
+    assert_eq!(result, "");
+}
+
+/// Test that tojson filter accepts ensure_ascii kwarg (HuggingFace compatibility)
+/// This is the fix for: "unknown keyword argument 'ensure_ascii'"
+#[test]
+fn test_tojson_with_ensure_ascii() {
+    // Template that uses tojson(ensure_ascii=False) like HuggingFace templates do
+    let template = r#"
+{%- for message in messages -%}
+{{ message.role }}: {{ message.content }}
+{%- if message.tool_calls is defined and message.tool_calls -%}
+Tools: {{ message.tool_calls|tojson(ensure_ascii=False) }}
+{%- endif -%}
+{% endfor -%}
+"#;
+
+    let processor = ChatTemplateProcessor::new(template.to_string());
+
+    let messages = [ChatMessage::User {
+        content: MessageContent::Text("Test with Unicode: 日本語".to_string()),
+        name: None,
+    }];
+
+    // Convert to JSON values
+    let json_messages: Vec<serde_json::Value> = messages
+        .iter()
+        .map(|msg| serde_json::to_value(msg).unwrap())
+        .collect();
+
+    // This should NOT fail with "unknown keyword argument 'ensure_ascii'"
+    let result = processor
+        .apply_chat_template(&json_messages, ChatTemplateParams::default())
+        .unwrap();
+
+    assert!(result.contains("user: Test with Unicode: 日本語"));
+}
+
+/// Test tojson with all HuggingFace kwargs
+#[test]
+fn test_tojson_with_all_huggingface_kwargs() {
+    // Template using all the kwargs that HuggingFace's custom tojson accepts
+    let template = r#"
+{%- set data = {"z_key": 1, "a_key": 2, "m_key": 3} -%}
+Unsorted: {{ data|tojson }}
+Sorted: {{ data|tojson(sort_keys=True) }}
+Indented: {{ data|tojson(indent=2) }}
+All: {{ data|tojson(ensure_ascii=False, sort_keys=True, indent=2) }}
+"#;
+
+    let processor = ChatTemplateProcessor::new(template.to_string());
+    let messages: Vec<serde_json::Value> = vec![];
+
+    // This should NOT fail - all kwargs should be accepted
+    let result = processor
+        .apply_chat_template(&messages, ChatTemplateParams::default())
+        .unwrap();
+
+    // Verify sorted output contains keys in alphabetical order
+    assert!(result.contains("Sorted:"));
+    // The sorted output should have a_key before m_key before z_key
+    let sorted_line = result.lines().find(|l| l.starts_with("Sorted:")).unwrap();
+    let a_pos = sorted_line.find("a_key").unwrap();
+    let m_pos = sorted_line.find("m_key").unwrap();
+    let z_pos = sorted_line.find("z_key").unwrap();
+    assert!(a_pos < m_pos && m_pos < z_pos, "Keys should be sorted");
+
+    // Verify indented output is pretty-printed with newlines
+    assert!(
+        result.contains("Indented: {\n"),
+        "Indented JSON should be pretty-printed with newlines"
+    );
+}
+
+#[test]
+fn test_content_format_detection() {
+    let string_template = r#"
+{%- for message in messages -%}
+{{ message.role }}: {{ message.content }}
+{%- endfor -%}
+"#;
+    assert_eq!(
+        detect_chat_template_content_format(string_template),
+        ChatTemplateContentFormat::String
+    );
+
+    let openai_template = r#"
+{%- for message in messages -%}
+  {%- for content in message.content -%}
+    {{ content.type }}: {{ content.text }}
+  {%- endfor -%}
+{%- endfor -%}
+"#;
+    assert_eq!(
+        detect_chat_template_content_format(openai_template),
+        ChatTemplateContentFormat::OpenAI
+    );
+}
+
+#[test]
+fn test_template_with_multimodal_content() {
+    let template = r#"
+{%- for message in messages %}
+{{ message.role }}:
+{%- if message.content is string %}
+{{ message.content }}
+{%- else %}
+{%- for part in message.content %}
+  {%- if part.type == "text" %}
+{{ part.text }}
+  {%- elif part.type == "image_url" %}
+[IMAGE]
+  {%- endif %}
+{%- endfor %}
+{%- endif %}
+{% endfor %}
+"#;
+
+    let processor = ChatTemplateProcessor::new(template.to_string());
+
+    let messages = [ChatMessage::User {
+        content: MessageContent::Parts(vec![
+            ContentPart::Text {
+                text: "Look at this:".to_string(),
+            },
+            ContentPart::ImageUrl {
+                image_url: ImageUrl {
+                    url: "https://example.com/image.jpg".to_string(),
+                    detail: None,
+                },
+            },
+        ]),
+        name: None,
+    }];
+
+    // Convert to JSON values
+    let json_messages: Vec<serde_json::Value> = messages
+        .iter()
+        .map(|msg| serde_json::to_value(msg).unwrap())
+        .collect();
+
+    let result = processor
+        .apply_chat_template(&json_messages, ChatTemplateParams::default())
+        .unwrap();
+
+    // Should contain both text and image parts
+    assert!(result.contains("user:"));
+    assert!(result.contains("Look at this:"));
+    assert!(result.contains("[IMAGE]"));
+}

sgl-model-gateway/tests/chat_template_loading.rs

@@ -0,0 +1,230 @@
+#[cfg(test)]
+mod tests {
+    use std::fs;
+
+    use sgl_model_gateway::{
+        protocols::chat::{ChatMessage, MessageContent},
+        tokenizer::{chat_template::ChatTemplateParams, huggingface::HuggingFaceTokenizer},
+    };
+    use tempfile::TempDir;
+
+    #[test]
+    fn test_load_chat_template_from_file() {
+        // Create temporary directory
+        let temp_dir = TempDir::new().unwrap();
+        let template_path = temp_dir.path().join("template.jinja");
+
+        // Write a test template
+        let template_content = r#"
+{%- for message in messages %}
+    {{- '<|' + message['role'] + '|>' + message['content'] }}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|assistant|>' }}
+{%- endif %}
+"#;
+        fs::write(&template_path, template_content).unwrap();
+
+        // Create a mock tokenizer config
+        let tokenizer_config = r#"{
+            "version": "1.0",
+            "truncation": null,
+            "padding": null,
+            "added_tokens": [],
+            "normalizer": null,
+            "pre_tokenizer": {
+                "type": "Whitespace"
+            },
+            "post_processor": null,
+            "decoder": null,
+            "model": {
+                "type": "BPE",
+                "vocab": {
+                    "hello": 0,
+                    "world": 1,
+                    "<s>": 2,
+                    "</s>": 3
+                },
+                "merges": []
+            }
+        }"#;
+
+        let tokenizer_path = temp_dir.path().join("tokenizer.json");
+        fs::write(&tokenizer_path, tokenizer_config).unwrap();
+
+        // Load tokenizer with custom chat template
+        let tokenizer = HuggingFaceTokenizer::from_file_with_chat_template(
+            tokenizer_path.to_str().unwrap(),
+            Some(template_path.to_str().unwrap()),
+        )
+        .unwrap();
+
+        let messages = [
+            ChatMessage::User {
+                content: MessageContent::Text("Hello".to_string()),
+                name: None,
+            },
+            ChatMessage::Assistant {
+                content: Some(MessageContent::Text("Hi there".to_string())),
+                name: None,
+                tool_calls: None,
+                reasoning_content: None,
+            },
+        ];
+
+        // Convert to JSON values like the router does
+        let json_messages: Vec<serde_json::Value> = messages
+            .iter()
+            .map(|msg| serde_json::to_value(msg).unwrap())
+            .collect();
+
+        use sgl_model_gateway::tokenizer::chat_template::ChatTemplateParams;
+        let params = ChatTemplateParams {
+            add_generation_prompt: true,
+            ..Default::default()
+        };
+        let result = tokenizer
+            .apply_chat_template(&json_messages, params)
+            .unwrap();
+
+        assert!(result.contains("<|user|>Hello"));
+        assert!(result.contains("<|assistant|>Hi there"));
+        assert!(result.ends_with("<|assistant|>"));
+    }
+
+    #[test]
+    fn test_override_existing_template() {
+        // Create temporary directory
+        let temp_dir = TempDir::new().unwrap();
+
+        // Create tokenizer config with a built-in template
+        let tokenizer_config_path = temp_dir.path().join("tokenizer_config.json");
+        let config_with_template = r#"{
+            "chat_template": "built-in: {% for msg in messages %}{{ msg.content }}{% endfor %}"
+        }"#;
+        fs::write(&tokenizer_config_path, config_with_template).unwrap();
+
+        // Create the actual tokenizer file
+        let tokenizer_json = r#"{
+            "version": "1.0",
+            "truncation": null,
+            "padding": null,
+            "added_tokens": [],
+            "normalizer": null,
+            "pre_tokenizer": {
+                "type": "Whitespace"
+            },
+            "post_processor": null,
+            "decoder": null,
+            "model": {
+                "type": "BPE",
+                "vocab": {
+                    "test": 0,
+                    "<s>": 1,
+                    "</s>": 2
+                },
+                "merges": []
+            }
+        }"#;
+        let tokenizer_path = temp_dir.path().join("tokenizer.json");
+        fs::write(&tokenizer_path, tokenizer_json).unwrap();
+
+        // Create custom template that should override
+        let custom_template_path = temp_dir.path().join("custom.jinja");
+        let custom_template =
+            r#"CUSTOM: {% for msg in messages %}[{{ msg.role }}]: {{ msg.content }}{% endfor %}"#;
+        fs::write(&custom_template_path, custom_template).unwrap();
+
+        // Load with custom template - should override the built-in one
+        let tokenizer = HuggingFaceTokenizer::from_file_with_chat_template(
+            tokenizer_path.to_str().unwrap(),
+            Some(custom_template_path.to_str().unwrap()),
+        )
+        .unwrap();
+
+        let messages = [ChatMessage::User {
+            content: MessageContent::Text("Test".to_string()),
+            name: None,
+        }];
+
+        // Convert to JSON values
+        let json_messages: Vec<serde_json::Value> = messages
+            .iter()
+            .map(|msg| serde_json::to_value(msg).unwrap())
+            .collect();
+
+        let result = tokenizer
+            .apply_chat_template(&json_messages, ChatTemplateParams::default())
+            .unwrap();
+
+        // Should use CUSTOM template, not built-in
+        assert!(result.starts_with("CUSTOM:"));
+        assert!(result.contains("[user]: Test"));
+        assert!(!result.contains("built-in:"));
+    }
+
+    #[test]
+    fn test_set_chat_template_after_creation() {
+        // Create temporary directory and tokenizer file
+        let temp_dir = TempDir::new().unwrap();
+        let tokenizer_json = r#"{
+            "version": "1.0",
+            "truncation": null,
+            "padding": null,
+            "added_tokens": [],
+            "normalizer": null,
+            "pre_tokenizer": {
+                "type": "Whitespace"
+            },
+            "post_processor": null,
+            "decoder": null,
+            "model": {
+                "type": "BPE",
+                "vocab": {
+                    "test": 0,
+                    "<s>": 1,
+                    "</s>": 2
+                },
+                "merges": []
+            }
+        }"#;
+        let tokenizer_path = temp_dir.path().join("tokenizer.json");
+        fs::write(&tokenizer_path, tokenizer_json).unwrap();
+
+        // Load tokenizer without custom template
+        let mut tokenizer =
+            HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap()).unwrap();
+
+        // Set a template after creation (mimics Python's behavior)
+        let new_template =
+            "NEW: {% for msg in messages %}{{ msg.role }}: {{ msg.content }}; {% endfor %}";
+        tokenizer.set_chat_template(new_template.to_string());
+
+        let messages = [
+            ChatMessage::User {
+                content: MessageContent::Text("Hello".to_string()),
+                name: None,
+            },
+            ChatMessage::Assistant {
+                content: Some(MessageContent::Text("World".to_string())),
+                name: None,
+                tool_calls: None,
+                reasoning_content: None,
+            },
+        ];
+
+        // Convert to JSON values
+        let json_messages: Vec<serde_json::Value> = messages
+            .iter()
+            .map(|msg| serde_json::to_value(msg).unwrap())
+            .collect();
+
+        let result = tokenizer
+            .apply_chat_template(&json_messages, ChatTemplateParams::default())
+            .unwrap();
+
+        assert!(result.starts_with("NEW:"));
+        assert!(result.contains("user: Hello;"));
+        assert!(result.contains("assistant: World;"));
+    }
+}

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

@@ -0,0 +1,174 @@
+// 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>,
+}
+
+#[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;
+    }
+}

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()
+}

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

@@ -0,0 +1,659 @@
+// These modules are used by tests and benchmarks
+#![allow(dead_code)]
+
+pub mod mock_mcp_server;
+pub mod mock_openai_server;
+pub mod mock_worker;
+pub mod streaming_helpers;
+pub mod test_app;
+
+use std::{
+    fs,
+    path::PathBuf,
+    sync::{Arc, Mutex, OnceLock},
+};
+
+use serde_json::json;
+use sgl_model_gateway::{
+    app_context::AppContext,
+    config::{RouterConfig, RoutingMode},
+    core::{
+        BasicWorkerBuilder, LoadMonitor, ModelCard, RuntimeType, Worker, WorkerRegistry, WorkerType,
+    },
+    data_connector::{
+        MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage,
+    },
+    middleware::TokenBucket,
+    policies::PolicyRegistry,
+    protocols::common::{Function, Tool},
+};
+
+/// Helper function to create AppContext for tests
+pub async fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
+    let client = reqwest::Client::new();
+
+    // Initialize rate limiter
+    let rate_limiter = match config.max_concurrent_requests {
+        n if n <= 0 => None,
+        n => {
+            let rate_limit_tokens = 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(config.policy.clone()));
+
+    // Initialize storage backends (Memory for tests)
+    let response_storage = Arc::new(MemoryResponseStorage::new());
+    let conversation_storage = Arc::new(MemoryConversationStorage::new());
+    let conversation_item_storage = Arc::new(MemoryConversationItemStorage::new());
+
+    // Initialize load monitor
+    let load_monitor = Some(Arc::new(LoadMonitor::new(
+        worker_registry.clone(),
+        policy_registry.clone(),
+        client.clone(),
+        config.worker_startup_check_interval_secs,
+    )));
+
+    // Create empty OnceLock for worker job queue, workflow engine, and mcp manager
+    let worker_job_queue = Arc::new(OnceLock::new());
+    let workflow_engine = Arc::new(OnceLock::new());
+    let mcp_manager_lock = Arc::new(OnceLock::new());
+
+    let app_context = Arc::new(
+        AppContext::builder()
+            .router_config(config.clone())
+            .client(client)
+            .rate_limiter(rate_limiter)
+            .tokenizer(None) // tokenizer
+            .reasoning_parser_factory(None) // reasoning_parser_factory
+            .tool_parser_factory(None) // tool_parser_factory
+            .worker_registry(worker_registry)
+            .policy_registry(policy_registry)
+            .response_storage(response_storage)
+            .conversation_storage(conversation_storage)
+            .conversation_item_storage(conversation_item_storage)
+            .load_monitor(load_monitor)
+            .worker_job_queue(worker_job_queue)
+            .workflow_engine(workflow_engine)
+            .mcp_manager(mcp_manager_lock)
+            .build()
+            .unwrap(),
+    );
+
+    // Initialize JobQueue after AppContext is created
+    let weak_context = Arc::downgrade(&app_context);
+    let job_queue = sgl_model_gateway::core::JobQueue::new(
+        sgl_model_gateway::core::JobQueueConfig::default(),
+        weak_context,
+    );
+    app_context
+        .worker_job_queue
+        .set(job_queue)
+        .expect("JobQueue should only be initialized once");
+
+    // Initialize WorkflowEngine and register workflows
+    use sgl_model_gateway::core::workflow::{
+        create_worker_registration_workflow, create_worker_removal_workflow, WorkflowEngine,
+    };
+    let engine = Arc::new(WorkflowEngine::new());
+    engine.register_workflow(create_worker_registration_workflow(&config));
+    engine.register_workflow(create_worker_removal_workflow());
+    app_context
+        .workflow_engine
+        .set(engine)
+        .expect("WorkflowEngine should only be initialized once");
+
+    // Register external workers for OpenAI mode
+    if let RoutingMode::OpenAI { worker_urls, .. } = &config.mode {
+        for url in worker_urls {
+            // Create a worker that supports common test models
+            let models = vec![
+                ModelCard::new("mock-model"),
+                ModelCard::new("gpt-4"),
+                ModelCard::new("gpt-3.5-turbo"),
+            ];
+            let worker: Arc<dyn Worker> = Arc::new(
+                BasicWorkerBuilder::new(url)
+                    .worker_type(WorkerType::Regular)
+                    .runtime_type(RuntimeType::External)
+                    .models(models)
+                    .build(),
+            );
+            app_context.worker_registry.register(worker);
+        }
+    }
+
+    // Initialize MCP manager with empty config
+    use sgl_model_gateway::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
+}
+
+/// Helper function to create AppContext for tests with MCP config from file
+pub async fn create_test_context_with_mcp_config(
+    config: RouterConfig,
+    mcp_config_path: &str,
+) -> Arc<AppContext> {
+    use sgl_model_gateway::mcp::{McpConfig, McpManager};
+
+    let client = reqwest::Client::new();
+
+    // Initialize rate limiter
+    let rate_limiter = match config.max_concurrent_requests {
+        n if n <= 0 => None,
+        n => {
+            let rate_limit_tokens = 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(config.policy.clone()));
+
+    // Initialize storage backends (Memory for tests)
+    let response_storage = Arc::new(MemoryResponseStorage::new());
+    let conversation_storage = Arc::new(MemoryConversationStorage::new());
+    let conversation_item_storage = Arc::new(MemoryConversationItemStorage::new());
+
+    // Initialize load monitor
+    let load_monitor = Some(Arc::new(LoadMonitor::new(
+        worker_registry.clone(),
+        policy_registry.clone(),
+        client.clone(),
+        config.worker_startup_check_interval_secs,
+    )));
+
+    // Create empty OnceLock for worker job queue, workflow engine, and mcp manager
+    let worker_job_queue = Arc::new(OnceLock::new());
+    let workflow_engine = Arc::new(OnceLock::new());
+    let mcp_manager_lock = Arc::new(OnceLock::new());
+
+    let app_context = Arc::new(
+        AppContext::builder()
+            .router_config(config.clone())
+            .client(client)
+            .rate_limiter(rate_limiter)
+            .tokenizer(None) // tokenizer
+            .reasoning_parser_factory(None) // reasoning_parser_factory
+            .tool_parser_factory(None) // tool_parser_factory
+            .worker_registry(worker_registry)
+            .policy_registry(policy_registry)
+            .response_storage(response_storage)
+            .conversation_storage(conversation_storage)
+            .conversation_item_storage(conversation_item_storage)
+            .load_monitor(load_monitor)
+            .worker_job_queue(worker_job_queue)
+            .workflow_engine(workflow_engine)
+            .mcp_manager(mcp_manager_lock)
+            .build()
+            .unwrap(),
+    );
+
+    // Initialize JobQueue after AppContext is created
+    let weak_context = Arc::downgrade(&app_context);
+    let job_queue = sgl_model_gateway::core::JobQueue::new(
+        sgl_model_gateway::core::JobQueueConfig::default(),
+        weak_context,
+    );
+    app_context
+        .worker_job_queue
+        .set(job_queue)
+        .expect("JobQueue should only be initialized once");
+
+    // Initialize WorkflowEngine and register workflows
+    use sgl_model_gateway::core::workflow::{
+        create_worker_registration_workflow, create_worker_removal_workflow, WorkflowEngine,
+    };
+    let engine = Arc::new(WorkflowEngine::new());
+    engine.register_workflow(create_worker_registration_workflow(&config));
+    engine.register_workflow(create_worker_removal_workflow());
+    app_context
+        .workflow_engine
+        .set(engine)
+        .expect("WorkflowEngine should only be initialized once");
+
+    // Register external workers for OpenAI mode
+    if let RoutingMode::OpenAI { worker_urls, .. } = &config.mode {
+        for url in worker_urls {
+            // Create a worker that supports common test models
+            let models = vec![
+                ModelCard::new("mock-model"),
+                ModelCard::new("gpt-4"),
+                ModelCard::new("gpt-3.5-turbo"),
+            ];
+            let worker: Arc<dyn Worker> = Arc::new(
+                BasicWorkerBuilder::new(url)
+                    .worker_type(WorkerType::Regular)
+                    .runtime_type(RuntimeType::External)
+                    .models(models)
+                    .build(),
+            );
+            app_context.worker_registry.register(worker);
+        }
+    }
+
+    // Initialize MCP manager from config file
+    let mcp_config = McpConfig::from_file(mcp_config_path)
+        .await
+        .expect("Failed to load MCP config from file");
+    let mcp_manager = McpManager::with_defaults(mcp_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
+}
+
+// Tokenizer download configuration
+const TINYLLAMA_TOKENIZER_URL: &str =
+    "https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/tokenizer.json";
+const CACHE_DIR: &str = ".tokenizer_cache";
+const TINYLLAMA_TOKENIZER_FILENAME: &str = "tinyllama_tokenizer.json";
+
+// Global mutex to prevent concurrent downloads
+static DOWNLOAD_MUTEX: OnceLock<Mutex<()>> = OnceLock::new();
+
+/// Downloads the TinyLlama tokenizer from HuggingFace if not already cached.
+/// Returns the path to the cached tokenizer file.
+///
+/// This function is thread-safe and will only download the tokenizer once
+/// even if called from multiple threads concurrently.
+pub fn ensure_tokenizer_cached() -> PathBuf {
+    // Get or initialize the mutex
+    let mutex = DOWNLOAD_MUTEX.get_or_init(|| Mutex::new(()));
+
+    // Lock to ensure only one thread downloads at a time
+    let _guard = mutex.lock().unwrap();
+
+    let cache_dir = PathBuf::from(CACHE_DIR);
+    let tokenizer_path = cache_dir.join(TINYLLAMA_TOKENIZER_FILENAME);
+
+    // Create cache directory if it doesn't exist
+    if !cache_dir.exists() {
+        fs::create_dir_all(&cache_dir).expect("Failed to create cache directory");
+    }
+
+    // Download tokenizer if not already cached
+    if !tokenizer_path.exists() {
+        println!("Downloading TinyLlama tokenizer from HuggingFace...");
+
+        // Use blocking reqwest client since we're in tests/benchmarks
+        let client = reqwest::blocking::Client::new();
+        let response = client
+            .get(TINYLLAMA_TOKENIZER_URL)
+            .send()
+            .expect("Failed to download tokenizer");
+
+        if !response.status().is_success() {
+            panic!("Failed to download tokenizer: HTTP {}", response.status());
+        }
+
+        let content = response.bytes().expect("Failed to read tokenizer content");
+
+        if content.len() < 100 {
+            panic!("Downloaded content too small: {} bytes", content.len());
+        }
+
+        fs::write(&tokenizer_path, content).expect("Failed to write tokenizer to cache");
+        println!(
+            "Tokenizer downloaded and cached successfully ({} bytes)",
+            tokenizer_path.metadata().unwrap().len()
+        );
+    }
+
+    tokenizer_path
+}
+
+/// Common test prompts for consistency across tests
+pub const TEST_PROMPTS: [&str; 4] = [
+    "deep learning is",
+    "Deep learning is",
+    "has anyone seen nemo lately",
+    "another prompt",
+];
+
+/// Pre-computed hashes for verification
+pub const EXPECTED_HASHES: [u64; 4] = [
+    1209591529327510910,
+    4181375434596349981,
+    6245658446118930933,
+    5097285695902185237,
+];
+
+/// Create a comprehensive set of test tools covering all parser test scenarios
+#[allow(dead_code)]
+pub fn create_test_tools() -> Vec<Tool> {
+    vec![
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "search".to_string(),
+                description: Some("Search for information".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "query": {"type": "string"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "get_weather".to_string(),
+                description: Some("Get weather information".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "city": {"type": "string"},
+                        "location": {"type": "string"},
+                        "date": {"type": "string"},
+                        "units": {"type": "string"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "calculate".to_string(),
+                description: Some("Perform calculations".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "x": {"type": "number"},
+                        "y": {"type": "number"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "translate".to_string(),
+                description: Some("Translate text".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "text": {"type": "string"},
+                        "to": {"type": "string"},
+                        "target_lang": {"type": "string"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "get_time".to_string(),
+                description: Some("Get current time".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "timezone": {"type": "string"},
+                        "format": {"type": "string"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "get_current_time".to_string(),
+                description: Some("Get current time".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "timezone": {"type": "string"},
+                        "format": {"type": "string"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "update_settings".to_string(),
+                description: Some("Update settings".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "preferences": {"type": "object"},
+                        "notifications": {"type": "boolean"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "ping".to_string(),
+                description: Some("Ping service".to_string()),
+                parameters: json!({"type": "object", "properties": {}}),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "test".to_string(),
+                description: Some("Test function".to_string()),
+                parameters: json!({"type": "object", "properties": {}}),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "process".to_string(),
+                description: Some("Process data".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "count": {"type": "number"},
+                        "rate": {"type": "number"},
+                        "enabled": {"type": "boolean"},
+                        "data": {"type": "object"},
+                        "text": {"type": "string"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "web_search".to_string(),
+                description: Some("Search the web".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "query": {"type": "string"},
+                        "num_results": {"type": "number"},
+                        "search_type": {"type": "string"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "get_tourist_attractions".to_string(),
+                description: Some("Get tourist attractions".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "city": {"type": "string"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "config".to_string(),
+                description: Some("Configuration function".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "debug": {"type": "boolean"},
+                        "verbose": {"type": "boolean"},
+                        "optional": {"type": "null"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "test_func".to_string(),
+                description: Some("Test function".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "bool_true": {"type": "boolean"},
+                        "bool_false": {"type": "boolean"},
+                        "none_val": {"type": "null"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "create".to_string(),
+                description: Some("Create resource".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "name": {"type": "string"},
+                        "email": {"type": "string"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "add".to_string(),
+                description: Some("Add operation".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "x": {"type": "number"},
+                        "y": {"type": "number"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "calc".to_string(),
+                description: Some("Calculate".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "x": {"type": "number"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "func1".to_string(),
+                description: Some("Function 1".to_string()),
+                parameters: json!({"type": "object", "properties": {}}),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "func2".to_string(),
+                description: Some("Function 2".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "y": {"type": "number"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "tool1".to_string(),
+                description: Some("Tool 1".to_string()),
+                parameters: json!({"type": "object", "properties": {}}),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: Function {
+                name: "tool2".to_string(),
+                description: Some("Tool 2".to_string()),
+                parameters: json!({
+                    "type": "object",
+                    "properties": {
+                        "y": {"type": "number"}
+                    }
+                }),
+                strict: None,
+            },
+        },
+    ]
+}

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 "
+    }
+}

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

@@ -0,0 +1,260 @@
+use std::sync::{Arc, OnceLock};
+
+use axum::Router;
+use reqwest::Client;
+use sgl_model_gateway::{
+    app_context::AppContext,
+    config::RouterConfig,
+    core::{
+        BasicWorkerBuilder, LoadMonitor, ModelCard, RuntimeType, Worker, WorkerRegistry, WorkerType,
+    },
+    data_connector::{
+        MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage,
+    },
+    mcp::{McpConfig, McpManager},
+    middleware::{AuthConfig, TokenBucket},
+    policies::PolicyRegistry,
+    routers::RouterTrait,
+    server::{build_app, AppState},
+};
+
+/// 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 engine
+    let worker_job_queue = Arc::new(OnceLock::new());
+    let workflow_engine = 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(None) // tokenizer
+            .reasoning_parser_factory(None) // reasoning_parser_factory
+            .tool_parser_factory(None) // tool_parser_factory
+            .worker_registry(worker_registry)
+            .policy_registry(policy_registry)
+            .response_storage(response_storage)
+            .conversation_storage(conversation_storage)
+            .conversation_item_storage(conversation_item_storage)
+            .load_monitor(load_monitor)
+            .worker_job_queue(worker_job_queue)
+            .workflow_engine(workflow_engine)
+            .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,
+    });
+
+    // 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,
+        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,
+    });
+
+    // 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,
+        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_engine = 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(None)
+            .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_engine(workflow_engine)
+            .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);
+}

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 sgl_model_gateway::mcp::{McpConfig, McpError, 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);
+}

sgl-model-gateway/tests/metrics_aggregator_test.rs

@@ -0,0 +1,271 @@
+use sgl_model_gateway::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);
+}

sgl-model-gateway/tests/multimodal_tracker_test.rs

@@ -0,0 +1,151 @@
+use std::{collections::HashMap, path::PathBuf, sync::Arc, time::Duration};
+
+use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
+use reqwest::Client;
+use sgl_model_gateway::multimodal::{
+    AsyncMultiModalTracker, ChatContentPart, ConversationSegment, ImageFetchConfig, ImageSource,
+    MediaConnector, MediaConnectorConfig, MediaSource, Modality, TrackerConfig,
+};
+use tempfile::tempdir;
+
+const TINY_PNG_BASE64: &str =
+    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNgYAAAAAMAASsJTYQAAAAASUVORK5CYII=";
+
+fn tiny_png_bytes() -> Vec<u8> {
+    BASE64_STANDARD
+        .decode(TINY_PNG_BASE64)
+        .expect("decode tiny png fixture")
+}
+
+fn test_connector(allowed_path: Option<PathBuf>) -> MediaConnector {
+    let client = Client::builder()
+        .timeout(Duration::from_secs(5))
+        .no_proxy()
+        .build()
+        .expect("client");
+    MediaConnector::new(
+        client,
+        MediaConnectorConfig {
+            allowed_domains: None,
+            allowed_local_media_path: allowed_path,
+            fetch_timeout: Duration::from_secs(5),
+        },
+    )
+    .expect("media connector")
+}
+
+#[tokio::test]
+async fn fetch_image_from_inline_bytes() {
+    let connector = test_connector(None);
+    let bytes = tiny_png_bytes();
+    let frame = connector
+        .fetch_image(
+            MediaSource::InlineBytes(bytes.clone()),
+            ImageFetchConfig::default(),
+        )
+        .await
+        .expect("inline image");
+    assert_eq!(frame.data().width(), 1);
+    assert_eq!(frame.data().height(), 1);
+    assert_eq!(frame.raw_bytes(), bytes.as_slice());
+}
+
+#[tokio::test]
+async fn fetch_image_from_data_url() {
+    let connector = test_connector(None);
+    let bytes = tiny_png_bytes();
+    let data_url = format!(
+        "data:image/png;base64,{}",
+        BASE64_STANDARD.encode(bytes.clone())
+    );
+
+    let frame = connector
+        .fetch_image(MediaSource::DataUrl(data_url), ImageFetchConfig::default())
+        .await
+        .expect("data url");
+    assert_eq!(frame.data().width(), 1);
+    assert_eq!(frame.raw_bytes(), bytes.as_slice());
+}
+
+#[tokio::test]
+async fn fetch_image_from_file() {
+    let tmp = tempdir().expect("tempdir");
+    let allowed_root = std::fs::canonicalize(tmp.path()).expect("canonical tmp path");
+    let file_path = allowed_root.join("tiny.png");
+    std::fs::write(&file_path, tiny_png_bytes()).expect("write png");
+
+    let connector = test_connector(Some(allowed_root));
+    let frame = connector
+        .fetch_image(
+            MediaSource::File(file_path.clone()),
+            ImageFetchConfig::default(),
+        )
+        .await
+        .expect("file png");
+    assert_eq!(frame.data().width(), 1);
+    let expected = std::fs::canonicalize(&file_path).expect("canonical path");
+    match frame.source() {
+        ImageSource::File { path } => assert_eq!(path, &expected),
+        other => panic!("expected file source, got {:?}", other),
+    }
+}
+
+#[tokio::test]
+async fn tracker_collects_conversation_and_placeholders() {
+    let connector = Arc::new(test_connector(None));
+    let mut tracker = AsyncMultiModalTracker::new(
+        connector,
+        TrackerConfig {
+            placeholder_tokens: Default::default(),
+            modality_limits: HashMap::from([(Modality::Image, 2)]),
+        },
+    );
+
+    tracker
+        .push_part(ChatContentPart::Text {
+            text: "before".into(),
+        })
+        .expect("text part");
+    tracker
+        .push_part(ChatContentPart::ImageData {
+            data: tiny_png_bytes(),
+            mime_type: Some("image/png".into()),
+            uuid: Some("img-1".into()),
+            detail: None,
+        })
+        .expect("image part");
+    tracker
+        .push_part(ChatContentPart::Text {
+            text: "after".into(),
+        })
+        .expect("text part");
+
+    let output = tracker.finalize().await.expect("tracker finalize");
+    assert_eq!(output.conversation.len(), 3);
+    assert!(matches!(
+        &output.conversation[0],
+        ConversationSegment::Text(text) if text == "before"
+    ));
+    assert!(matches!(
+        &output.conversation[1],
+        ConversationSegment::Placeholder { token } if token == "<image>"
+    ));
+    assert!(matches!(
+        &output.conversation[2],
+        ConversationSegment::Text(text) if text == "after"
+    ));
+
+    let images = output.data.get(&Modality::Image).expect("image entry");
+    assert_eq!(images.len(), 1);
+
+    let uuids = output.uuids.get(&Modality::Image).expect("uuid entry");
+    assert_eq!(uuids, &vec![Some("img-1".into())]);
+
+    let placeholders = output
+        .placeholders
+        .get("<image>")
+        .expect("placeholder entry");
+    assert_eq!(placeholders.len(), 1);
+    assert_eq!(placeholders[0].text_position, 1);
+    assert_eq!(placeholders[0].item_index, 0);
+}

sgl-model-gateway/tests/policy_registry_integration.rs

@@ -0,0 +1,170 @@
+//! Integration tests for PolicyRegistry with RouterManager
+
+use std::{collections::HashMap, sync::Arc};
+
+use sgl_model_gateway::{
+    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,
+        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,
+        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,
+        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 sgl_model_gateway::{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 sgl_model_gateway::{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");
+}

sgl-model-gateway/tests/request_formats_test.rs

@@ -0,0 +1,358 @@
+mod common;
+
+use std::sync::Arc;
+
+use common::mock_worker::{HealthStatus, MockWorker, MockWorkerConfig, WorkerType};
+use reqwest::Client;
+use serde_json::json;
+use sgl_model_gateway::{
+    config::{RouterConfig, RoutingMode},
+    routers::{RouterFactory, RouterTrait},
+};
+
+/// Test context that manages mock workers
+struct TestContext {
+    workers: Vec<MockWorker>,
+    _router: Arc<dyn RouterTrait>,
+    worker_urls: Vec<String>,
+}
+
+impl TestContext {
+    async fn new(worker_configs: Vec<MockWorkerConfig>) -> Self {
+        let mut config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .port(3003)
+            .worker_startup_timeout_secs(1)
+            .worker_startup_check_interval_secs(1)
+            .build_unchecked();
+
+        let mut workers = Vec::new();
+        let mut worker_urls = Vec::new();
+
+        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;
+        }
+
+        config.mode = RoutingMode::Regular {
+            worker_urls: worker_urls.clone(),
+        };
+
+        let app_context = common::create_test_context(config.clone()).await;
+
+        let router = RouterFactory::create_router(&app_context).await.unwrap();
+        let router = Arc::from(router);
+
+        if !workers.is_empty() {
+            tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
+        }
+
+        Self {
+            workers,
+            _router: router,
+            worker_urls: worker_urls.clone(),
+        }
+    }
+
+    async fn shutdown(mut self) {
+        // Small delay to ensure any pending operations complete
+        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+
+        for worker in &mut self.workers {
+            worker.stop().await;
+        }
+
+        // Another small delay to ensure cleanup completes
+        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+    }
+
+    async fn make_request(
+        &self,
+        endpoint: &str,
+        body: serde_json::Value,
+    ) -> Result<serde_json::Value, String> {
+        let client = Client::new();
+
+        // Use the first worker URL from the context
+        let worker_url = self
+            .worker_urls
+            .first()
+            .ok_or_else(|| "No workers available".to_string())?;
+
+        let response = client
+            .post(format!("{}{}", worker_url, endpoint))
+            .json(&body)
+            .send()
+            .await
+            .map_err(|e| format!("Request failed: {}", e))?;
+
+        if !response.status().is_success() {
+            return Err(format!("Request failed with status: {}", response.status()));
+        }
+
+        response
+            .json::<serde_json::Value>()
+            .await
+            .map_err(|e| format!("Failed to parse response: {}", e))
+    }
+}
+
+#[cfg(test)]
+mod request_format_tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_generate_request_formats() {
+        let ctx = TestContext::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 = TestContext::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 = TestContext::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 = TestContext::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 = TestContext::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 = TestContext::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;
+    }
+}

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

@@ -0,0 +1,570 @@
+use serde_json::json;
+use sgl_model_gateway::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
+    );
+}

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

@@ -0,0 +1,91 @@
+use serde_json::json;
+use sgl_model_gateway::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");
+}

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

@@ -0,0 +1,96 @@
+use serde_json::{from_str, json, to_string};
+use sgl_model_gateway::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()),
+    };
+
+    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,
+    };
+
+    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,
+    };
+    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,
+    };
+    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,
+    };
+    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,
+    };
+    // Only top-level string elements are extracted
+    assert_eq!(req.extract_text_for_routing(), "a");
+}

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;

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 sgl_model_gateway::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);
+}

sgl-model-gateway/tests/spec_test.rs

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

sgl-model-gateway/tests/streaming_tests.rs

@@ -0,0 +1,347 @@
+mod common;
+
+use std::sync::Arc;
+
+use common::mock_worker::{HealthStatus, MockWorker, MockWorkerConfig, WorkerType};
+use futures_util::StreamExt;
+use reqwest::Client;
+use serde_json::json;
+use sgl_model_gateway::{
+    config::{RouterConfig, RoutingMode},
+    routers::{RouterFactory, RouterTrait},
+};
+
+/// Test context that manages mock workers
+struct TestContext {
+    workers: Vec<MockWorker>,
+    _router: Arc<dyn RouterTrait>,
+    worker_urls: Vec<String>,
+}
+
+impl TestContext {
+    async fn new(worker_configs: Vec<MockWorkerConfig>) -> Self {
+        let mut config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .port(3004)
+            .worker_startup_timeout_secs(1)
+            .worker_startup_check_interval_secs(1)
+            .build_unchecked();
+
+        let mut workers = Vec::new();
+        let mut worker_urls = Vec::new();
+
+        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;
+        }
+
+        config.mode = RoutingMode::Regular {
+            worker_urls: worker_urls.clone(),
+        };
+
+        let app_context = common::create_test_context(config.clone()).await;
+
+        let router = RouterFactory::create_router(&app_context).await.unwrap();
+        let router = Arc::from(router);
+
+        if !workers.is_empty() {
+            tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
+        }
+
+        Self {
+            workers,
+            _router: router,
+            worker_urls: worker_urls.clone(),
+        }
+    }
+
+    async fn shutdown(mut self) {
+        // Small delay to ensure any pending operations complete
+        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+
+        for worker in &mut self.workers {
+            worker.stop().await;
+        }
+
+        // Another small delay to ensure cleanup completes
+        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+    }
+
+    async fn make_streaming_request(
+        &self,
+        endpoint: &str,
+        body: serde_json::Value,
+    ) -> Result<Vec<String>, String> {
+        let client = Client::new();
+
+        // Use the first worker URL from the context
+        let worker_url = self
+            .worker_urls
+            .first()
+            .ok_or_else(|| "No workers available".to_string())?;
+
+        let response = client
+            .post(format!("{}{}", worker_url, endpoint))
+            .json(&body)
+            .send()
+            .await
+            .map_err(|e| format!("Request failed: {}", e))?;
+
+        if !response.status().is_success() {
+            return Err(format!("Request failed with status: {}", response.status()));
+        }
+
+        // Check if it's a streaming response
+        let content_type = response
+            .headers()
+            .get("content-type")
+            .and_then(|v| v.to_str().ok())
+            .unwrap_or("");
+
+        if !content_type.contains("text/event-stream") {
+            return Err("Response is not a stream".to_string());
+        }
+
+        let mut stream = response.bytes_stream();
+        let mut events = Vec::new();
+
+        while let Some(chunk) = stream.next().await {
+            if let Ok(bytes) = chunk {
+                let text = String::from_utf8_lossy(&bytes);
+                for line in text.lines() {
+                    if let Some(stripped) = line.strip_prefix("data: ") {
+                        events.push(stripped.to_string());
+                    }
+                }
+            }
+        }
+
+        Ok(events)
+    }
+}
+
+#[cfg(test)]
+mod streaming_tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_generate_streaming() {
+        let ctx = TestContext::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();
+        // Should have at least one data chunk and [DONE]
+        assert!(events.len() >= 2);
+        assert_eq!(events.last().unwrap(), "[DONE]");
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_v1_chat_completions_streaming() {
+        let ctx = TestContext::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); // At least one chunk + [DONE]
+
+        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 = TestContext::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); // At least one chunk + [DONE]
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_streaming_with_error() {
+        let ctx = TestContext::new(vec![MockWorkerConfig {
+            port: 20004,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 0,
+            fail_rate: 1.0, // Always fail
+        }])
+        .await;
+
+        let payload = json!({
+            "text": "This should fail",
+            "stream": true
+        });
+
+        let result = ctx.make_streaming_request("/generate", payload).await;
+        // With fail_rate: 1.0, the request should fail
+        assert!(result.is_err());
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_streaming_timeouts() {
+        let ctx = TestContext::new(vec![MockWorkerConfig {
+            port: 20005,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 100, // Slow response
+            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();
+
+        // Should have received multiple chunks over time
+        assert!(!events.is_empty());
+        assert!(elapsed.as_millis() >= 100); // At least one delay
+
+        ctx.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn test_batch_streaming() {
+        let ctx = TestContext::new(vec![MockWorkerConfig {
+            port: 20006,
+            worker_type: WorkerType::Regular,
+            health_status: HealthStatus::Healthy,
+            response_delay_ms: 10,
+            fail_rate: 0.0,
+        }])
+        .await;
+
+        // Batch request with streaming
+        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();
+        // Should have multiple events for batch
+        assert!(events.len() >= 4); // At least 3 responses + [DONE]
+
+        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]");
+    }
+}

sgl-model-gateway/tests/test_openai_routing.rs

@@ -0,0 +1,905 @@
+//! Comprehensive integration tests for OpenAI backend functionality
+
+use std::{
+    collections::HashMap,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+};
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{Method, StatusCode},
+    response::Response,
+    routing::post,
+    Json, Router,
+};
+use serde_json::json;
+use sgl_model_gateway::{
+    config::{
+        ConfigError, ConfigValidator, HistoryBackend, OracleConfig, RouterConfig, RoutingMode,
+    },
+    data_connector::{ResponseId, StoredResponse},
+    protocols::{
+        chat::{ChatCompletionRequest, ChatMessage, MessageContent},
+        common::StringOrArray,
+        completion::CompletionRequest,
+        generate::GenerateRequest,
+        responses::{ResponseInput, ResponsesGetParams, ResponsesRequest},
+    },
+    routers::{openai::OpenAIRouter, RouterTrait},
+};
+use tokio::{
+    net::TcpListener,
+    time::{sleep, Duration},
+};
+use tower::ServiceExt;
+
+mod common;
+use 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 = common::test_app::create_test_app_context().await;
+    // Register an external worker before creating the router
+    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 = common::test_app::create_test_app_context().await;
+    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 = common::test_app::create_test_app_context().await;
+    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 = common::test_app::create_test_app_context().await;
+    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 = common::test_app::create_test_app_context().await;
+    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,
+        sgl_model_gateway::config::PolicyConfig::Random,
+    );
+
+    let app_context = common::create_test_context(router_config).await;
+
+    let router = sgl_model_gateway::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 = common::test_app::create_test_app_context().await;
+    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 = common::test_app::create_test_app_context().await;
+    // Register the mock server worker and create router
+    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 = common::test_app::create_test_app_context().await;
+    // Register the mock server worker and create router
+    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 = common::test_app::create_test_app_context().await;
+    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 = common::test_app::create_test_app_context().await;
+    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 = common::test_app::create_test_app_context().await;
+    // Register a worker with the default model
+    common::test_app::register_external_worker(&ctx, "https://api.example.com", None);
+    let router = OpenAIRouter::new(&ctx).await.unwrap();
+
+    // Get models - should return the registered model
+    let req = Request::builder()
+        .method(Method::GET)
+        .uri("/models")
+        .body(Body::empty())
+        .unwrap();
+
+    let response = router.get_models(req).await;
+    assert_eq!(response.status(), StatusCode::OK);
+    let (_, body) = response.into_parts();
+    let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap();
+    let body_str = String::from_utf8(body_bytes.to_vec()).unwrap();
+    let models: serde_json::Value = serde_json::from_str(&body_str).unwrap();
+    assert_eq!(models["object"], "list");
+
+    // Should have the default model (gpt-3.5-turbo)
+    let data = models["data"].as_array().unwrap();
+    assert_eq!(data.len(), 1);
+    assert_eq!(data[0]["id"], "gpt-3.5-turbo");
+}
+
+#[test]
+fn oracle_config_validation_requires_config_when_enabled() {
+    let config = RouterConfig::builder()
+        .openai_mode(vec!["https://api.openai.com".to_string()])
+        .history_backend(HistoryBackend::Oracle)
+        .build_unchecked();
+
+    let err =
+        ConfigValidator::validate(&config).expect_err("config should fail without oracle details");
+
+    match err {
+        ConfigError::MissingRequired { field } => {
+            assert_eq!(field, "oracle");
+        }
+        other => panic!("unexpected error: {:?}", other),
+    }
+}
+
+#[test]
+fn oracle_config_validation_accepts_dsn_only() {
+    let config = RouterConfig::builder()
+        .openai_mode(vec!["https://api.openai.com".to_string()])
+        .oracle_history(OracleConfig {
+            wallet_path: None,
+            connect_descriptor: "tcps://db.example.com:1522/service".to_string(),
+            username: "scott".to_string(),
+            password: "tiger".to_string(),
+            pool_min: 1,
+            pool_max: 4,
+            pool_timeout_secs: 30,
+        })
+        .build_unchecked();
+
+    ConfigValidator::validate(&config).expect("dsn-based config should validate");
+}
+
+#[test]
+fn oracle_config_validation_accepts_wallet_alias() {
+    let config = RouterConfig::builder()
+        .openai_mode(vec!["https://api.openai.com".to_string()])
+        .oracle_history(OracleConfig {
+            wallet_path: Some("/etc/sglang/oracle-wallet".to_string()),
+            connect_descriptor: "db_low".to_string(),
+            username: "app_user".to_string(),
+            password: "secret".to_string(),
+            pool_min: 1,
+            pool_max: 8,
+            pool_timeout_secs: 45,
+        })
+        .build_unchecked();
+
+    ConfigValidator::validate(&config).expect("wallet-based config should validate");
+}

sgl-model-gateway/tests/test_pd_routing.rs

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

sgl-model-gateway/tests/tokenizer_cache_correctness_test.rs

@@ -0,0 +1,467 @@
+//! Cache correctness integration test
+//!
+//! This test validates that the tokenizer cache (L0, L1, and L0+L1 combined) produces
+//! exactly the same token IDs as uncached tokenization across multiple chat turns.
+//! Uses the real Qwen/Qwen3-4B-Instruct-2507 tokenizer to test with actual special tokens.
+
+use std::{
+    path::PathBuf,
+    sync::{Arc, OnceLock},
+};
+
+use sgl_model_gateway::tokenizer::{
+    cache::{CacheConfig, CachedTokenizer},
+    hub::download_tokenizer_from_hf,
+    huggingface::HuggingFaceTokenizer,
+    traits::Encoder,
+};
+
+/// Global tokenizer path cache - download once, reuse across all tests
+static TOKENIZER_PATH: OnceLock<Option<PathBuf>> = OnceLock::new();
+
+/// Download Qwen3-4B-Instruct-2507 tokenizer once and cache the path
+async fn get_tokenizer_path() -> Option<PathBuf> {
+    // Check if already downloaded
+    if let Some(cached) = TOKENIZER_PATH.get() {
+        return cached.clone();
+    }
+
+    // Download tokenizer
+    let result = match download_tokenizer_from_hf("Qwen/Qwen3-4B-Instruct-2507").await {
+        Ok(cache_dir) => {
+            let tokenizer_path = cache_dir.join("tokenizer.json");
+            if tokenizer_path.exists() {
+                Some(tokenizer_path)
+            } else {
+                println!("Tokenizer downloaded but tokenizer.json not found");
+                None
+            }
+        }
+        Err(e) => {
+            println!("Failed to download tokenizer: {}", e);
+            None
+        }
+    };
+
+    // Cache the result (even if None, so we don't retry on failure)
+    TOKENIZER_PATH.set(result.clone()).ok();
+    result
+}
+
+/// Comprehensive multi-turn chat conversation for testing cache correctness
+/// Uses Qwen's special tokens with diverse content to hit edge cases
+const CHAT_TURNS: [&str; 29] = [
+    // Basic conversation patterns
+    "<|im_start|>system\nYou are a helpful AI assistant.<|im_end|>",
+    "<|im_start|>system\nYou are a helpful AI assistant.<|im_end|><|im_start|>user\nWhat is the capital of France?<|im_end|>",
+    "<|im_start|>system\nYou are a helpful AI assistant.<|im_end|><|im_start|>user\nWhat is the capital of France?<|im_end|><|im_start|>assistant\nThe capital of France is Paris.<|im_end|>",
+
+    // Different system prompts (testing different prefix patterns)
+    "<|im_start|>system\nYou are a coding tutor specializing in Rust programming.<|im_end|><|im_start|>user\nExplain ownership.<|im_end|>",
+    "<|im_start|>system\nYou are a math teacher.<|im_end|><|im_start|>user\nSolve: 2x + 5 = 13<|im_end|>",
+
+    // Long conversation with multiple turns (testing longer prefixes)
+    "<|im_start|>system\nYou are a helpful AI assistant.<|im_end|><|im_start|>user\nTell me about deep learning.<|im_end|><|im_start|>assistant\nDeep learning is a subset of machine learning that uses neural networks with multiple layers.<|im_end|><|im_start|>user\nWhat are the main architectures?<|im_end|>",
+
+    // Code snippets (testing different character patterns)
+    "<|im_start|>system\nYou are a code reviewer.<|im_end|><|im_start|>user\nReview this code:\nfn main() {\n    println!(\"Hello, world!\");\n}\n<|im_end|>",
+    "<|im_start|>system\nYou are a code reviewer.<|im_end|><|im_start|>user\nExplain this Rust code:\nimpl<T> Drop for Box<T> {\n    fn drop(&mut self) { /* ... */ }\n}\n<|im_end|>",
+
+    // Mathematical content
+    "<|im_start|>system\nYou are a math tutor.<|im_end|><|im_start|>user\nProve that √2 is irrational using proof by contradiction.<|im_end|>",
+    "<|im_start|>system\nYou are a math tutor.<|im_end|><|im_start|>user\nCalculate: ∫(x² + 3x + 2)dx from 0 to 5<|im_end|>",
+
+    // Multilingual content
+    "<|im_start|>system\nYou are a multilingual assistant.<|im_end|><|im_start|>user\nTranslate to French: The quick brown fox jumps over the lazy dog.<|im_end|>",
+    "<|im_start|>system\nYou are a multilingual assistant.<|im_end|><|im_start|>user\n你好，请帮我翻译这句话：I love programming in Rust.<|im_end|>",
+    "<|im_start|>system\nYou are a multilingual assistant.<|im_end|><|im_start|>user\nこんにちは！Rustについて教えてください。<|im_end|>",
+
+    // Special characters and emojis
+    "<|im_start|>system\nYou are a friendly chatbot.<|im_end|><|im_start|>user\nWhat do you think about emojis? 😀🎉🚀💻<|im_end|>",
+    "<|im_start|>system\nYou are a data analyst.<|im_end|><|im_start|>user\nAnalyze this: {\"name\": \"test\", \"value\": 42, \"nested\": {\"key\": \"value\"}}<|im_end|>",
+
+    // Very long message (testing large token counts)
+    "<|im_start|>system\nYou are a literature expert.<|im_end|><|im_start|>user\nAnalyze the themes in this passage: In the vast expanse of the digital realm, where bits and bytes dance in harmonious symphony, there exists a paradigm that transcends mere computation. This paradigm, known as machine learning, represents humanity's quest to imbue silicon with the spark of cognition. Deep neural networks, inspired by the intricate architecture of biological brains, layer upon layer of artificial neurons, each connection a synapse firing in the dark recesses of mathematical space. Through gradient descent, these networks learn patterns invisible to human perception, extracting meaning from chaos, signal from noise. The transformer architecture revolutionized this field, introducing attention mechanisms that allowed models to focus on relevant information, much like how humans selectively attend to important details in their environment.<|im_end|>",
+
+    // Edge case: Multiple special tokens in sequence
+    "<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nHi<|im_end|><|im_start|>assistant\nHello!<|im_end|><|im_start|>user\nHow are you?<|im_end|>",
+
+    // Edge case: Empty-ish messages
+    "<|im_start|>system\n<|im_end|><|im_start|>user\nTest<|im_end|>",
+    "<|im_start|>system\nBrief.<|im_end|><|im_start|>user\nOK<|im_end|>",
+
+    // Technical documentation style
+    "<|im_start|>system\nYou are a technical writer.<|im_end|><|im_start|>user\nDocument the following API:\n\n```rust\npub struct CachedTokenizer {\n    inner: Arc<dyn Tokenizer>,\n    l0: Option<L0Cache>,\n    l1: Option<L1Cache>,\n}\n\nimpl Encoder for CachedTokenizer {\n    fn encode(&self, input: &str) -> Result<Encoding>;\n}\n```\n<|im_end|>",
+
+    // Conversation with code review
+    "<|im_start|>system\nYou are a senior Rust developer.<|im_end|><|im_start|>user\nReview for correctness:\n\nlet special_tokens: Option<Vec<&str>> = self.l1.as_ref().map(|_| {\n    self.special_token_strings.iter().map(|s| s.as_str()).collect()\n});<|im_end|>",
+
+    // Markdown formatted content
+    "<|im_start|>system\nYou are a documentation assistant.<|im_end|><|im_start|>user\nFormat this as markdown:\n\n# Cache Architecture\n\n## L0 Cache\n- Exact match\n- DashMap based\n- 10K entries\n\n## L1 Cache  \n- Prefix match\n- Special token boundaries\n- 50MB memory\n<|im_end|>",
+
+    // Complex nested structures
+    "<|im_start|>system\nYou are a JSON expert.<|im_end|><|im_start|>user\nValidate this JSON:\n{\n  \"tokenizer_cache\": {\n    \"enable_l0\": true,\n    \"l0_max_entries\": 10000,\n    \"enable_l1\": true,\n    \"l1_max_memory\": 52428800,\n    \"stats\": {\n      \"hits\": [1, 2, 3],\n      \"misses\": {\"count\": 5}\n    }\n  }\n}\n<|im_end|>",
+
+    // SQL queries
+    "<|im_start|>system\nYou are a database expert.<|im_end|><|im_start|>user\nOptimize this query:\nSELECT u.name, COUNT(p.id) as post_count\nFROM users u\nLEFT JOIN posts p ON u.id = p.user_id\nWHERE u.created_at > '2024-01-01'\nGROUP BY u.id, u.name\nHAVING COUNT(p.id) > 5\nORDER BY post_count DESC;<|im_end|>",
+
+    // Regex patterns
+    "<|im_start|>system\nYou are a regex expert.<|im_end|><|im_start|>user\nExplain this regex: ^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\\.)+[a-zA-Z]{2,}$<|im_end|>",
+
+    // Command line examples
+    "<|im_start|>system\nYou are a DevOps engineer.<|im_end|><|im_start|>user\nExplain this command:\ncargo bench --bench tokenizer_benchmark -- --color=never | tee results.txt<|im_end|>",
+
+    // Unicode edge cases
+    "<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nTest: café, naïve, Zürich, 北京, 東京, मुंबई, Москва<|im_end|>",
+
+    // Mixed content complexity
+    "<|im_start|>system\nYou are a software architect.<|im_end|><|im_start|>user\nDesign a caching system that:\n1. Handles 10K+ QPS\n2. Maintains 99.9% uptime  \n3. Supports L0 (exact) and L1 (prefix) caching\n4. Uses Blake3 for hashing (10GB/s throughput)\n5. Implements LRU eviction\n6. Thread-safe with lock-free reads\n\nKey requirements:\n- Memory: 50MB L1 budget\n- Latency: <100µs p99\n- Correctness: 100% (no false tokens)\n<|im_end|>",
+
+    // Very long technical discussion
+    "<|im_start|>system\nYou are a compiler expert.<|im_end|><|im_start|>user\nExplain why BPE tokenizers are not prefix-stable:\n\nThe core issue is that BPE applies merges based on local context. When you tokenize 'prefix' alone, it might apply merge rules differently than when tokenizing 'prefix + suffix' as a whole. For example:\n\ntokenize('hello world') might produce [hello, _world]\ntokenize('hello') + tokenize(' world') might produce [hel, lo, _wo, rld]\n\nThis is because the merge rules see different contexts. The space before 'world' in the first case is part of the token boundary, but in the second case, ' world' is tokenized in isolation.\n\nSpecial tokens solve this because they are:\n1. Atomic (never split or merged)\n2. Protected from normalization\n3. Marked with special: true flag\n4. Have normalized: false property\n\nThis guarantees: tokenize(prefix + special + suffix) = tokenize(prefix + special) + tokenize(suffix)\n\nOur L1 cache exploits this by:\n1. Finding all special token boundaries\n2. Re-tokenizing prefixes at those boundaries\n3. Caching the exact token IDs\n4. On cache hit, appending suffix tokens\n\nThis achieves both correctness (100%) and performance (22.7x speedup on high prefix reuse workloads).<|im_end|>",
+];
+
+#[tokio::test]
+async fn test_cache_produces_identical_tokens() {
+    // Get tokenizer path (download once, cached across tests)
+    let tokenizer_path = match get_tokenizer_path().await {
+        Some(path) => path,
+        None => {
+            println!("Skipping test - tokenizer not available");
+            return;
+        }
+    };
+
+    // Create base tokenizer (no cache)
+    let base_tokenizer = Arc::new(
+        HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
+            .expect("Failed to load base tokenizer"),
+    );
+
+    // Create cached tokenizers with different configurations
+    let l0_only_config = CacheConfig {
+        enable_l0: true,
+        l0_max_entries: 10_000,
+        enable_l1: false,
+        l1_max_memory: 0,
+    };
+
+    let l1_only_config = CacheConfig {
+        enable_l0: false,
+        l0_max_entries: 0,
+        enable_l1: true,
+        l1_max_memory: 50 * 1024 * 1024,
+    };
+
+    let l0_l1_config = CacheConfig {
+        enable_l0: true,
+        l0_max_entries: 10_000,
+        enable_l1: true,
+        l1_max_memory: 50 * 1024 * 1024,
+    };
+
+    let l0_tokenizer = Arc::new(CachedTokenizer::new(base_tokenizer.clone(), l0_only_config));
+    let l1_tokenizer = Arc::new(CachedTokenizer::new(base_tokenizer.clone(), l1_only_config));
+    let l0_l1_tokenizer = Arc::new(CachedTokenizer::new(base_tokenizer.clone(), l0_l1_config));
+
+    println!(
+        "\n=== Testing Cache Correctness Across {} Chat Turns ===\n",
+        CHAT_TURNS.len()
+    );
+
+    for (turn_idx, turn) in CHAT_TURNS.iter().enumerate() {
+        println!("Turn {}: Testing {} chars", turn_idx + 1, turn.len());
+
+        // Tokenize with base (no cache)
+        let base_encoding = base_tokenizer
+            .encode(turn)
+            .expect("Base tokenization failed");
+        let base_tokens = base_encoding.token_ids();
+
+        // Tokenize with L0-only
+        let l0_encoding = l0_tokenizer.encode(turn).expect("L0 tokenization failed");
+        let l0_tokens = l0_encoding.token_ids();
+
+        // Tokenize with L1-only
+        let l1_encoding = l1_tokenizer.encode(turn).expect("L1 tokenization failed");
+        let l1_tokens = l1_encoding.token_ids();
+
+        // Tokenize with L0+L1
+        let l0_l1_encoding = l0_l1_tokenizer
+            .encode(turn)
+            .expect("L0+L1 tokenization failed");
+        let l0_l1_tokens = l0_l1_encoding.token_ids();
+
+        // Verify all configurations produce identical token IDs
+        assert_eq!(
+            base_tokens.len(),
+            l0_tokens.len(),
+            "Turn {}: L0 token count mismatch (base: {}, L0: {})",
+            turn_idx + 1,
+            base_tokens.len(),
+            l0_tokens.len()
+        );
+
+        assert_eq!(
+            base_tokens.len(),
+            l1_tokens.len(),
+            "Turn {}: L1 token count mismatch (base: {}, L1: {})",
+            turn_idx + 1,
+            base_tokens.len(),
+            l1_tokens.len()
+        );
+
+        assert_eq!(
+            base_tokens.len(),
+            l0_l1_tokens.len(),
+            "Turn {}: L0+L1 token count mismatch (base: {}, L0+L1: {})",
+            turn_idx + 1,
+            base_tokens.len(),
+            l0_l1_tokens.len()
+        );
+
+        // Compare token by token
+        for (token_idx, (((base_token, l0_token), l1_token), l0_l1_token)) in base_tokens
+            .iter()
+            .zip(l0_tokens.iter())
+            .zip(l1_tokens.iter())
+            .zip(l0_l1_tokens.iter())
+            .enumerate()
+        {
+            assert_eq!(
+                base_token,
+                l0_token,
+                "Turn {}, token {}: L0 mismatch (base: {}, L0: {})",
+                turn_idx + 1,
+                token_idx,
+                base_token,
+                l0_token
+            );
+
+            assert_eq!(
+                base_token,
+                l1_token,
+                "Turn {}, token {}: L1 mismatch (base: {}, L1: {})",
+                turn_idx + 1,
+                token_idx,
+                base_token,
+                l1_token
+            );
+
+            assert_eq!(
+                base_token,
+                l0_l1_token,
+                "Turn {}, token {}: L0+L1 mismatch (base: {}, L0+L1: {})",
+                turn_idx + 1,
+                token_idx,
+                base_token,
+                l0_l1_token
+            );
+        }
+
+        println!(
+            "  ✓ All configurations produced identical {} tokens",
+            base_tokens.len()
+        );
+    }
+
+    // Print cache statistics
+    if let Some(l0_stats) = l0_tokenizer.cache_stats() {
+        println!("\n=== L0 Cache Statistics ===");
+        println!("  Hits: {}", l0_stats.hits);
+        println!("  Misses: {}", l0_stats.misses);
+        println!(
+            "  Hit rate: {:.2}%",
+            if l0_stats.hits + l0_stats.misses > 0 {
+                l0_stats.hits as f64 / (l0_stats.hits + l0_stats.misses) as f64 * 100.0
+            } else {
+                0.0
+            }
+        );
+        println!("  Entries: {}", l0_stats.entries);
+    }
+
+    if let Some(l1_stats) = l1_tokenizer.l1_cache_stats() {
+        println!("\n=== L1 Cache Statistics ===");
+        println!("  Hits: {}", l1_stats.hits);
+        println!("  Misses: {}", l1_stats.misses);
+        println!(
+            "  Hit rate: {:.2}%",
+            if l1_stats.hits + l1_stats.misses > 0 {
+                l1_stats.hits as f64 / (l1_stats.hits + l1_stats.misses) as f64 * 100.0
+            } else {
+                0.0
+            }
+        );
+        println!("  Entries: {}", l1_stats.entries);
+        println!("  Memory used: {} bytes", l1_stats.memory_bytes);
+    }
+
+    if let Some(l0_stats) = l0_l1_tokenizer.cache_stats() {
+        if let Some(l1_stats) = l0_l1_tokenizer.l1_cache_stats() {
+            println!("\n=== L0+L1 Combined Cache Statistics ===");
+            println!("  L0 Hits: {}", l0_stats.hits);
+            println!("  L1 Hits: {}", l1_stats.hits);
+            println!(
+                "  Total Hit rate: {:.2}%",
+                if l0_stats.hits + l1_stats.hits + l0_stats.misses + l1_stats.misses > 0 {
+                    (l0_stats.hits + l1_stats.hits) as f64
+                        / (l0_stats.hits + l1_stats.hits + l0_stats.misses + l1_stats.misses) as f64
+                        * 100.0
+                } else {
+                    0.0
+                }
+            );
+        }
+    }
+
+    println!("\n✓ All cache configurations produce identical tokenization results!");
+}
+
+#[tokio::test]
+async fn test_cache_correctness_with_edge_cases() {
+    // Get tokenizer path (download once, cached across tests)
+    let tokenizer_path = match get_tokenizer_path().await {
+        Some(path) => path,
+        None => {
+            println!("Skipping test - tokenizer not available");
+            return;
+        }
+    };
+
+    // Create base and cached tokenizers
+    let base_tokenizer = Arc::new(
+        HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
+            .expect("Failed to load base tokenizer"),
+    );
+
+    let cached_config = CacheConfig {
+        enable_l0: true,
+        l0_max_entries: 10_000,
+        enable_l1: true,
+        l1_max_memory: 50 * 1024 * 1024,
+    };
+
+    let cached_tokenizer = Arc::new(CachedTokenizer::new(base_tokenizer.clone(), cached_config));
+
+    println!("\n=== Testing Edge Cases and Complex Patterns ===\n");
+
+    // Edge cases that stress-test the cache
+    let edge_cases = [
+        // Minimal messages
+        ("<|im_start|>system\n<|im_end|>", "Empty system message"),
+        ("<|im_start|>user\na<|im_end|>", "Single character"),
+
+        // Special token boundaries
+        ("<|im_start|>system\nA<|im_end|><|im_start|>user\nB<|im_end|><|im_start|>assistant\nC<|im_end|>", "Minimal multi-turn"),
+
+        // Repeated exact queries (L0 hit test)
+        ("<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nHello!<|im_end|>", "Repeated query 1"),
+        ("<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nHello!<|im_end|>", "Repeated query 2"),
+
+        // Same prefix, different suffix (L1 hit test)
+        ("<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nWhat is 1+1?<|im_end|>", "Same prefix, diff suffix 1"),
+        ("<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nWhat is 2+2?<|im_end|>", "Same prefix, diff suffix 2"),
+        ("<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nWhat is 3+3?<|im_end|>", "Same prefix, diff suffix 3"),
+
+        // Unicode stress tests
+        ("<|im_start|>system\n你好<|im_end|><|im_start|>user\n世界<|im_end|>", "Chinese characters"),
+        ("<|im_start|>system\nこんにちは<|im_end|><|im_start|>user\n世界<|im_end|>", "Japanese + Chinese"),
+        ("<|im_start|>system\n🚀💻🎉<|im_end|><|im_start|>user\n😀😃😄<|im_end|>", "Emoji only"),
+
+        // Whitespace edge cases
+        ("<|im_start|>system\n   \n<|im_end|>", "Whitespace only"),
+        ("<|im_start|>system\n\n\n\n<|im_end|>", "Multiple newlines"),
+        ("<|im_start|>system\n\t\t\t<|im_end|>", "Tabs"),
+
+        // Long token sequences
+        ("<|im_start|>system\nThe quick brown fox jumps over the lazy dog. The quick brown fox jumps over the lazy dog. The quick brown fox jumps over the lazy dog.<|im_end|>", "Repeated phrase"),
+
+        // Special characters
+        ("<|im_start|>system\n!@#$%^&*()_+-=[]{}|;':\",./<>?<|im_end|>", "ASCII special chars"),
+        ("<|im_start|>system\n`~\\<|im_end|>", "Backtick and tilde"),
+
+        // Code with special formatting
+        ("<|im_start|>system\nCode: fn() -> Result<(), Box<dyn Error>><|im_end|>", "Rust generics"),
+        ("<|im_start|>system\nRegex: ^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}$<|im_end|>", "Email regex"),
+
+        // Very long single token sequences (testing buffer handling)
+        ("<|im_start|>system\naaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa<|im_end|>", "Repeated 'a'"),
+        ("<|im_start|>system\n0123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789<|im_end|>", "Repeated numbers"),
+    ];
+
+    let mut test_count = 0;
+    let mut mismatch_count = 0;
+
+    for (query, description) in edge_cases.iter() {
+        test_count += 1;
+
+        let base_tokens = base_tokenizer
+            .encode(query)
+            .expect("Base encoding failed")
+            .token_ids()
+            .to_vec();
+
+        let cached_tokens = cached_tokenizer
+            .encode(query)
+            .expect("Cached encoding failed")
+            .token_ids()
+            .to_vec();
+
+        if base_tokens != cached_tokens {
+            mismatch_count += 1;
+            println!("  ✗ {}: Token mismatch!", description);
+            println!(
+                "    Base length: {}, Cached length: {}",
+                base_tokens.len(),
+                cached_tokens.len()
+            );
+
+            // Show first few mismatching tokens for debugging
+            for (i, (base, cached)) in base_tokens.iter().zip(cached_tokens.iter()).enumerate() {
+                if base != cached {
+                    println!("    Token {}: base={}, cached={}", i, base, cached);
+                    if i >= 5 {
+                        break;
+                    }
+                }
+            }
+        } else {
+            println!("  ✓ {}: {} tokens", description, base_tokens.len());
+        }
+    }
+
+    assert_eq!(
+        mismatch_count, 0,
+        "{} out of {} edge cases failed!",
+        mismatch_count, test_count
+    );
+
+    // Print cache statistics
+    if let Some(l0_stats) = cached_tokenizer.cache_stats() {
+        println!("\n=== Cache Statistics ===");
+        println!(
+            "  L0 Hits: {} ({:.1}% hit rate)",
+            l0_stats.hits,
+            if l0_stats.hits + l0_stats.misses > 0 {
+                l0_stats.hits as f64 / (l0_stats.hits + l0_stats.misses) as f64 * 100.0
+            } else {
+                0.0
+            }
+        );
+    }
+
+    if let Some(l1_stats) = cached_tokenizer.l1_cache_stats() {
+        println!(
+            "  L1 Hits: {} ({:.1}% hit rate)",
+            l1_stats.hits,
+            if l1_stats.hits + l1_stats.misses > 0 {
+                l1_stats.hits as f64 / (l1_stats.hits + l1_stats.misses) as f64 * 100.0
+            } else {
+                0.0
+            }
+        );
+    }
+
+    println!("\n✓ All {} edge cases passed!", test_count);
+}

sgl-model-gateway/tests/tokenizer_integration.rs

@@ -0,0 +1,560 @@
+//! Integration tests for tokenizers using real tokenizer data
+//!
+//! These tests download the TinyLlama tokenizer from HuggingFace to verify our tokenizer
+//! implementation works correctly with real-world tokenizer files.
+
+mod common;
+use std::sync::Arc;
+
+use common::{ensure_tokenizer_cached, EXPECTED_HASHES, TEST_PROMPTS};
+use sgl_model_gateway::tokenizer::{
+    factory, huggingface::HuggingFaceTokenizer, sequence::Sequence, stop::*, stream::DecodeStream,
+    traits::*,
+};
+
+const LONG_TEST_PROMPTS: [(&str, &str); 6] = [
+    ("Tell me about the following text.", "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat."),
+    ("Tell me about the following text.", "Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum."),
+    ("Tell me about the following text.", "Sed ut perspiciatis unde omnis iste natus error sit voluptatem accusantium doloremque laudantium, totam rem aperiam, eaque ipsa quae ab illo inventore veritatis et quasi architecto beatae vitae dicta sunt explicabo. Nemo enim ipsam voluptatem quia voluptas sit aspernatur aut odit aut fugit, sed quia consequuntur magni dolores eos qui ratione voluptatem sequi nesciunt."),
+    ("Tell me about the following text.", "Neque porro quisquam est, qui dolorem ipsum quia dolor sit amet, consectetur, adipisci velit, sed quia non numquam eius modi tempora incidunt ut labore et dolore magnam aliquam quaerat voluptatem."),
+    // Tennis-themed prompt for variety
+    ("Tell me about the following text.", "In the ancient realm of Tennisia, the very magic of the land is drawn from the sport itself. Forehands light the skies, backhands carve the earth, and serves rumble like thunder across kingdoms. At the center of this balance lie four sacred Grand Slam relics: the Sapphire Trophy of Melbourne, the Emerald Chalice of Paris, the Ruby Crown of London, and the Diamond Orb of New York. Together, they keep the game's spirit alive.
+    But the relics are scattered, guarded by champions of legendary skill. The first is the Fire King of Clay, ruler of the crimson courts, whose topspin arcs blaze high and heavy, scorching all who dare stand across from him. The second is the Tempest Trickster, master of the baseline fortress, whose footwork and precision can turn back any storm, and whose returns arrive as if pulled by invisible strings. The third is the Shadow-Dancer of the Highlands, a tactician who thrives in the long rallies of twilight, changing pace and spin until opponents lose their rhythm. The fourth and final guardian is a towering Diamond Titan, a net-charging colossus whose volleys shatter the air itself.
+    Into this arena of gods steps the Silver-Wristed Knight — a player of impossible grace, whose game is an art form. His quest: to claim each relic not for glory, but to restore harmony to the rankings of the realm.
+    He travels across the Kingdom of Clay, where the points stretch like marathons and the air tastes of iron; through the Grasslands of London, where the ball skids low and the margins are razor-thin; over the Hard Courts of the East, where rallies turn into duels of endurance; and finally to the Cathedral of Lights in New York, where night matches burn with fevered energy.
+    Each battle is played under enchanted floodlights, the lines patrolled by spectral line judges whose calls are final. The crowd's roar swells with every break point, and the Silver-Wristed Knight's racket glows brightest when the match teeters at deuce. There are moments when doubt grips him — when his serve falters or his touch deserts him — but each challenge teaches a new stroke, culminating in the legendary Forehand of Dawn.
+    When the last relic is claimed, he stands not as a conqueror but as a custodian of the game, knowing that rivalries forge the very magic he protects. The balance is restored — until the next season begins."),
+    // Emoji stress test
+    ("Tell me about the following text.", "😀😃😄😁😆🥹😅😂🤣🥲☺️😊😇🙂🙃😉🤩😎 🤪🥳🤓🙄🤪😵👻")
+];
+
+fn compute_hashes_for_tokenizer<E: Encoder>(tokenizer: &E, prompts: &[&str]) -> Vec<u64> {
+    prompts
+        .iter()
+        .map(|&prompt| {
+            tokenizer
+                .encode(prompt)
+                .expect("Failed to encode prompt")
+                .get_hash()
+        })
+        .collect()
+}
+
+#[test]
+fn test_huggingface_tokenizer_hashes() {
+    let tokenizer_path = ensure_tokenizer_cached();
+    let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
+        .expect("Failed to load HuggingFace tokenizer");
+
+    let prompt_hashes = compute_hashes_for_tokenizer(&tokenizer, &TEST_PROMPTS);
+
+    println!(
+        "HF Tokenizer: {:?}\nComputed Hashes: {:?}\nExpected Hashes: {:?}",
+        tokenizer_path, prompt_hashes, EXPECTED_HASHES
+    );
+
+    assert_eq!(prompt_hashes, EXPECTED_HASHES);
+}
+
+#[test]
+fn test_tokenizer_encode_decode_lifecycle() {
+    let tokenizer_path = ensure_tokenizer_cached();
+    let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
+        .expect("Failed to load HuggingFace tokenizer");
+
+    for prompt in TEST_PROMPTS.iter() {
+        let encoding = tokenizer.encode(prompt).expect("Failed to encode prompt");
+
+        let decoded = tokenizer
+            .decode(encoding.token_ids(), false)
+            .expect("Failed to decode token_ids");
+
+        assert_eq!(decoded, *prompt, "Encode-decode mismatch for: {}", prompt);
+    }
+}
+
+#[test]
+fn test_sequence_operations() {
+    let tokenizer_path = ensure_tokenizer_cached();
+    let tokenizer = Arc::new(
+        HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
+            .expect("Failed to load tokenizer"),
+    );
+
+    for prompt in TEST_PROMPTS.iter() {
+        let encoding = tokenizer.encode(prompt).expect("Failed to encode prompt");
+
+        let mut sequence = Sequence::new(tokenizer.clone());
+        sequence.append_text(prompt).expect("Failed to append text");
+
+        assert_eq!(
+            sequence.len(),
+            encoding.token_ids().len(),
+            "Sequence length mismatch"
+        );
+        assert_eq!(sequence.text().unwrap(), *prompt, "Sequence text mismatch");
+
+        let mut decoder = Sequence::new(tokenizer.clone());
+        let mut output = String::new();
+
+        for token_id in encoding.token_ids() {
+            let text = decoder
+                .append_token(*token_id)
+                .expect("Failed to append token");
+            output.push_str(&text);
+        }
+
+        assert_eq!(decoder.len(), sequence.len(), "Decoder length mismatch");
+        assert_eq!(
+            decoder.token_ids(),
+            sequence.token_ids(),
+            "Token IDs mismatch"
+        );
+        assert_eq!(output, *prompt, "Incremental decode mismatch");
+    }
+}
+
+#[test]
+fn test_decode_stream() {
+    let tokenizer_path = ensure_tokenizer_cached();
+    let tokenizer = Arc::new(
+        HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
+            .expect("Failed to load tokenizer"),
+    );
+
+    for prompt in TEST_PROMPTS.iter() {
+        let encoding = tokenizer.encode(prompt).expect("Failed to encode prompt");
+
+        let mut decoder = DecodeStream::new(tokenizer.clone(), &[], false);
+        let mut output = String::new();
+
+        for token_id in encoding.token_ids() {
+            if let Some(text) = decoder.step(*token_id).expect("Failed to decode token") {
+                output.push_str(&text);
+            }
+        }
+
+        assert_eq!(output, *prompt, "DecodeStream output mismatch");
+    }
+}
+
+#[test]
+fn test_long_sequence_incremental_decode_with_prefill() {
+    let tokenizer_path = ensure_tokenizer_cached();
+    let tokenizer = Arc::new(
+        HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
+            .expect("Failed to load tokenizer"),
+    );
+
+    for (input_text, output_text) in LONG_TEST_PROMPTS.iter() {
+        let input_encoding = tokenizer
+            .encode(input_text)
+            .expect("Failed to encode input");
+
+        let output_encoding = tokenizer
+            .encode(output_text)
+            .expect("Failed to encode output");
+
+        let mut decoder = DecodeStream::new(tokenizer.clone(), input_encoding.token_ids(), false);
+
+        let mut output = String::new();
+        for token_id in output_encoding.token_ids() {
+            if let Some(text) = decoder.step(*token_id).expect("Failed to decode token") {
+                output.push_str(&text);
+            }
+        }
+
+        assert_eq!(output.trim(), *output_text, "Long sequence decode mismatch");
+    }
+}
+
+#[test]
+fn test_stop_sequence_decoder() {
+    let tokenizer_path = ensure_tokenizer_cached();
+    let tokenizer = Arc::new(
+        HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
+            .expect("Failed to load tokenizer"),
+    );
+
+    let test_cases = vec![
+        (
+            "Hello world! Stop here. Continue after.",
+            "Stop",
+            "Hello world! ",
+        ),
+        ("Testing stop sequences.", ".", "Testing stop sequences"),
+        ("No stop sequence here", "xyz", "No stop sequence here"),
+    ];
+
+    for (input, stop_seq, expected) in test_cases {
+        let config = StopSequenceConfig::default().with_stop_sequence(stop_seq);
+
+        let mut decoder = StopSequenceDecoder::new(tokenizer.clone(), config, false);
+
+        let encoding = tokenizer.encode(input).expect("Failed to encode");
+        let mut output = String::new();
+        let mut stopped = false;
+
+        for token_id in encoding.token_ids() {
+            match decoder.process_token(*token_id).unwrap() {
+                SequenceDecoderOutput::Text(text) => output.push_str(&text),
+                SequenceDecoderOutput::StoppedWithText(text) => {
+                    output.push_str(&text);
+                    stopped = true;
+                    break;
+                }
+                SequenceDecoderOutput::Stopped => {
+                    stopped = true;
+                    break;
+                }
+                SequenceDecoderOutput::Held => {}
+            }
+        }
+
+        if !stopped {
+            // Flush any remaining text
+            if let SequenceDecoderOutput::Text(text) = decoder.flush() {
+                output.push_str(&text);
+            }
+        }
+
+        println!(
+            "Input: '{}', Stop: '{}', Output: '{}', Expected: '{}'",
+            input, stop_seq, output, expected
+        );
+
+        // The test should check if output starts with expected
+        // since stop sequences might not be perfectly aligned with token boundaries
+        assert!(
+            output.starts_with(expected) || output == input,
+            "Stop sequence test failed"
+        );
+    }
+}
+
+#[test]
+fn test_factory_creation() {
+    let tokenizer_path = ensure_tokenizer_cached();
+    let tokenizer = factory::create_tokenizer(tokenizer_path.to_str().unwrap())
+        .expect("Failed to create tokenizer via factory");
+
+    let encoding = tokenizer.encode(TEST_PROMPTS[0]).expect("Failed to encode");
+
+    let decoded = tokenizer
+        .decode(encoding.token_ids(), false)
+        .expect("Failed to decode");
+
+    assert_eq!(decoded, TEST_PROMPTS[0]);
+}
+
+#[test]
+fn test_batch_encoding() {
+    let tokenizer_path = ensure_tokenizer_cached();
+    let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
+        .expect("Failed to load tokenizer");
+
+    let encodings = tokenizer
+        .encode_batch(&TEST_PROMPTS)
+        .expect("Failed to batch encode");
+
+    assert_eq!(encodings.len(), TEST_PROMPTS.len());
+
+    for (i, encoding) in encodings.iter().enumerate() {
+        let decoded = tokenizer
+            .decode(encoding.token_ids(), false)
+            .expect("Failed to decode");
+        assert_eq!(decoded, TEST_PROMPTS[i]);
+    }
+}
+
+#[test]
+fn test_special_tokens() {
+    use sgl_model_gateway::tokenizer::traits::Tokenizer as TokenizerTrait;
+
+    let tokenizer_path = ensure_tokenizer_cached();
+    let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
+        .expect("Failed to load tokenizer");
+
+    let special_tokens = tokenizer.get_special_tokens();
+
+    // TinyLlama should have at least BOS and EOS tokens
+    assert!(special_tokens.bos_token.is_some());
+    assert!(special_tokens.eos_token.is_some());
+
+    println!("Special tokens: {:?}", special_tokens);
+}
+
+#[test]
+fn test_thread_safety() {
+    use std::thread;
+
+    let tokenizer_path = ensure_tokenizer_cached();
+    let tokenizer = Arc::new(
+        HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
+            .expect("Failed to load tokenizer"),
+    );
+
+    let handles: Vec<_> = TEST_PROMPTS
+        .iter()
+        .map(|&prompt| {
+            let tokenizer_clone = tokenizer.clone();
+            thread::spawn(move || {
+                let encoding = tokenizer_clone
+                    .encode(prompt)
+                    .expect("Failed to encode in thread");
+                let decoded = tokenizer_clone
+                    .decode(encoding.token_ids(), false)
+                    .expect("Failed to decode in thread");
+                assert_eq!(decoded, prompt);
+            })
+        })
+        .collect();
+
+    for handle in handles {
+        handle.join().expect("Thread panicked");
+    }
+}
+
+#[test]
+fn test_chat_template_discovery() {
+    use std::fs;
+
+    use tempfile::TempDir;
+
+    // Create a temporary directory with test files
+    let temp_dir = TempDir::new().expect("Failed to create temp dir");
+    let dir_path = temp_dir.path();
+
+    // Copy a real tokenizer.json file for testing
+    // We'll use the TinyLlama tokenizer that's already cached
+    let cached_tokenizer = ensure_tokenizer_cached();
+    let tokenizer_path = dir_path.join("tokenizer.json");
+    fs::copy(&cached_tokenizer, &tokenizer_path).expect("Failed to copy tokenizer file");
+
+    // Test 1: With chat_template.jinja file
+    let jinja_path = dir_path.join("chat_template.jinja");
+    fs::write(&jinja_path, "{{ messages }}").expect("Failed to write chat template");
+
+    let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap());
+    assert!(
+        tokenizer.is_ok(),
+        "Should load tokenizer with chat template"
+    );
+
+    // Clean up for next test
+    fs::remove_file(&jinja_path).ok();
+
+    // Test 2: With tokenizer_config.json containing chat_template
+    let config_path = dir_path.join("tokenizer_config.json");
+    fs::write(&config_path, r#"{"chat_template": "{{ messages }}"}"#)
+        .expect("Failed to write config");
+
+    let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap());
+    assert!(
+        tokenizer.is_ok(),
+        "Should load tokenizer with embedded template"
+    );
+
+    // Test 3: No chat template
+    fs::remove_file(&config_path).ok();
+    let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap());
+    assert!(
+        tokenizer.is_ok(),
+        "Should load tokenizer without chat template"
+    );
+}
+
+#[test]
+fn test_load_chat_template_from_local_file() {
+    use std::fs;
+
+    use tempfile::TempDir;
+
+    // Test 1: Load tokenizer with explicit chat template path
+    let temp_dir = TempDir::new().expect("Failed to create temp dir");
+    let dir_path = temp_dir.path();
+
+    // Copy a real tokenizer for testing
+    let cached_tokenizer = ensure_tokenizer_cached();
+    let tokenizer_path = dir_path.join("tokenizer.json");
+    fs::copy(&cached_tokenizer, &tokenizer_path).expect("Failed to copy tokenizer");
+
+    // Create a chat template file
+    let template_path = dir_path.join("my_template.jinja");
+    let template_content = r#"{% for message in messages %}{{ message.role }}: {{ message.content }}
+{% endfor %}"#;
+    fs::write(&template_path, template_content).expect("Failed to write template");
+
+    // Load tokenizer with explicit template path
+    let tokenizer = HuggingFaceTokenizer::from_file_with_chat_template(
+        tokenizer_path.to_str().unwrap(),
+        Some(template_path.to_str().unwrap()),
+    );
+    assert!(
+        tokenizer.is_ok(),
+        "Should load tokenizer with explicit template path"
+    );
+}
+
+#[tokio::test]
+async fn test_tinyllama_embedded_template() {
+    use sgl_model_gateway::tokenizer::hub::download_tokenizer_from_hf;
+
+    // Skip in CI without HF_TOKEN
+
+    // Test 2: TinyLlama has chat template embedded in tokenizer_config.json
+    match download_tokenizer_from_hf("TinyLlama/TinyLlama-1.1B-Chat-v1.0").await {
+        Ok(cache_dir) => {
+            // Verify tokenizer_config.json exists
+            let config_path = cache_dir.join("tokenizer_config.json");
+            assert!(config_path.exists(), "tokenizer_config.json should exist");
+
+            // Load the config and check for chat_template
+            let config_content =
+                std::fs::read_to_string(&config_path).expect("Failed to read config");
+            assert!(
+                config_content.contains("\"chat_template\""),
+                "TinyLlama should have embedded chat_template in config"
+            );
+
+            // Load tokenizer and verify it has chat template
+            let tokenizer_path = cache_dir.join("tokenizer.json");
+            let _tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
+                .expect("Failed to load tokenizer");
+
+            println!(
+                "✓ TinyLlama: Loaded tokenizer with embedded template from tokenizer_config.json"
+            );
+        }
+        Err(e) => {
+            println!("Download test skipped due to error: {}", e);
+        }
+    }
+}
+
+#[tokio::test]
+async fn test_qwen3_next_embedded_template() {
+    use sgl_model_gateway::tokenizer::hub::download_tokenizer_from_hf;
+
+    // Test 3: Qwen3-Next has chat template in tokenizer_config.json
+    match download_tokenizer_from_hf("Qwen/Qwen3-Next-80B-A3B-Instruct").await {
+        Ok(cache_dir) => {
+            let config_path = cache_dir.join("tokenizer_config.json");
+            assert!(config_path.exists(), "tokenizer_config.json should exist");
+
+            // Verify chat_template in config
+            let config_content =
+                std::fs::read_to_string(&config_path).expect("Failed to read config");
+            assert!(
+                config_content.contains("\"chat_template\""),
+                "Qwen3-Next should have chat_template in tokenizer_config.json"
+            );
+
+            // Load tokenizer
+            let tokenizer_path = cache_dir.join("tokenizer.json");
+            if tokenizer_path.exists() {
+                let _tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
+                    .expect("Failed to load tokenizer");
+                println!("✓ Qwen3-Next: Loaded tokenizer with embedded template");
+            }
+        }
+        Err(e) => {
+            println!("Download test skipped due to error: {}", e);
+        }
+    }
+}
+
+#[tokio::test]
+async fn test_qwen3_vl_json_template_priority() {
+    use sgl_model_gateway::tokenizer::hub::download_tokenizer_from_hf;
+
+    // Test 4: Qwen3-VL has both tokenizer_config.json template and chat_template.json
+    // Should prioritize chat_template.json
+    match download_tokenizer_from_hf("Qwen/Qwen3-VL-235B-A22B-Instruct").await {
+        Ok(cache_dir) => {
+            // Check for chat_template.json
+            let json_template_path = cache_dir.join("chat_template.json");
+            let has_json_template = json_template_path.exists();
+
+            // Also check tokenizer_config.json
+            let config_path = cache_dir.join("tokenizer_config.json");
+            assert!(config_path.exists(), "tokenizer_config.json should exist");
+
+            if has_json_template {
+                let json_content = std::fs::read_to_string(&json_template_path)
+                    .expect("Failed to read chat_template.json");
+                println!("✓ Qwen3-VL: Found chat_template.json (should be prioritized)");
+
+                // Verify it contains jinja template
+                assert!(
+                    !json_content.is_empty(),
+                    "chat_template.json should contain template"
+                );
+            }
+
+            // Load tokenizer - it should use the appropriate template
+            let tokenizer_path = cache_dir.join("tokenizer.json");
+            if tokenizer_path.exists() {
+                let _tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
+                    .expect("Failed to load tokenizer");
+                println!("✓ Qwen3-VL: Loaded tokenizer with template priority handling");
+            }
+        }
+        Err(e) => {
+            println!("Download test skipped due to error: {}", e);
+        }
+    }
+}
+
+#[tokio::test]
+async fn test_llava_separate_jinja_template() {
+    use sgl_model_gateway::tokenizer::hub::download_tokenizer_from_hf;
+
+    // Test 5: llava has chat_template.jinja as a separate file, not in tokenizer_config.json
+    match download_tokenizer_from_hf("llava-hf/llava-1.5-7b-hf").await {
+        Ok(cache_dir) => {
+            // Check for .jinja file
+            let jinja_path = cache_dir.join("chat_template.jinja");
+            let has_jinja = jinja_path.exists()
+                || std::fs::read_dir(&cache_dir)
+                    .map(|entries| {
+                        entries.filter_map(|e| e.ok()).any(|e| {
+                            e.file_name()
+                                .to_str()
+                                .is_some_and(|name| name.ends_with(".jinja"))
+                        })
+                    })
+                    .unwrap_or(false);
+
+            if has_jinja {
+                println!("✓ llava: Found separate .jinja chat template file");
+            }
+
+            // Check tokenizer_config.json - should NOT have embedded template
+            let config_path = cache_dir.join("tokenizer_config.json");
+            if config_path.exists() {
+                let config_content =
+                    std::fs::read_to_string(&config_path).expect("Failed to read config");
+
+                // llava might not have chat_template in config
+                if !config_content.contains("\"chat_template\"") {
+                    println!("✓ llava: No embedded template in config (as expected)");
+                }
+            }
+
+            // Load tokenizer - should auto-discover the .jinja file
+            let tokenizer_path = cache_dir.join("tokenizer.json");
+            if tokenizer_path.exists() {
+                let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap());
+                if tokenizer.is_ok() {
+                    println!("✓ llava: Loaded tokenizer with auto-discovered .jinja template");
+                } else {
+                    println!("Note: llava tokenizer loading failed - might need specific handling");
+                }
+            }
+        }
+        Err(e) => {
+            println!("Download test skipped due to error: {}", e);
+        }
+    }
+}

sgl-model-gateway/tests/tool_parser_deepseek.rs

@@ -0,0 +1,161 @@
+//! DeepSeek V3 Parser Integration Tests
+
+use sgl_model_gateway::tool_parser::{DeepSeekParser, ToolParser};
+
+mod common;
+use common::create_test_tools;
+
+#[tokio::test]
+async fn test_deepseek_complete_parsing() {
+    let parser = DeepSeekParser::new();
+
+    let input = r#"Let me help you with that.
+<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather
+```json
+{"location": "Tokyo", "units": "celsius"}
+```<｜tool▁call▁end｜><｜tool▁calls▁end｜>
+The weather in Tokyo is..."#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "Let me help you with that.\n");
+    assert_eq!(tools[0].function.name, "get_weather");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["location"], "Tokyo");
+    assert_eq!(args["units"], "celsius");
+}
+
+#[tokio::test]
+async fn test_deepseek_multiple_tools() {
+    let parser = DeepSeekParser::new();
+
+    let input = r#"<｜tool▁calls▁begin｜>
+<｜tool▁call▁begin｜>function<｜tool▁sep｜>search
+```json
+{"query": "rust programming"}
+```<｜tool▁call▁end｜>
+<｜tool▁call▁begin｜>function<｜tool▁sep｜>translate
+```json
+{"text": "Hello World", "to": "ja"}
+```<｜tool▁call▁end｜>
+<｜tool▁calls▁end｜>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 2);
+    assert_eq!(tools[0].function.name, "search");
+    assert_eq!(tools[1].function.name, "translate");
+}
+
+#[tokio::test]
+async fn test_deepseek_streaming() {
+    let tools = create_test_tools();
+
+    let mut parser = DeepSeekParser::new();
+
+    // Simulate streaming chunks
+    let chunks = vec![
+        "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>",
+        "function<｜tool▁sep｜>get_weather\n",
+        "```json\n",
+        r#"{"location": "#,
+        r#""Beijing", "#,
+        r#""units": "metric"}"#,
+        "\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>",
+    ];
+
+    let mut found_name = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+
+        for call in result.calls {
+            if let Some(name) = call.name {
+                assert_eq!(name, "get_weather");
+                found_name = true;
+            }
+        }
+    }
+
+    assert!(found_name, "Should have found tool name during streaming");
+}
+
+#[tokio::test]
+async fn test_deepseek_nested_json() {
+    let parser = DeepSeekParser::new();
+
+    let input = r#"<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>process
+```json
+{
+    "data": {
+        "nested": {
+            "deep": [1, 2, 3]
+        }
+    }
+}
+```<｜tool▁call▁end｜><｜tool▁calls▁end｜>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "process");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert!(args["data"]["nested"]["deep"].is_array());
+}
+
+#[test]
+fn test_deepseek_format_detection() {
+    let parser = DeepSeekParser::new();
+
+    // Should detect DeepSeek format
+    assert!(parser.has_tool_markers("<｜tool▁calls▁begin｜>"));
+    assert!(parser.has_tool_markers("text with <｜tool▁calls▁begin｜> marker"));
+
+    // Should not detect other formats
+    assert!(!parser.has_tool_markers("[TOOL_CALLS]"));
+    assert!(!parser.has_tool_markers("<tool_call>"));
+    assert!(!parser.has_tool_markers("plain text"));
+}
+
+#[tokio::test]
+async fn test_deepseek_malformed_json_handling() {
+    let parser = DeepSeekParser::new();
+
+    // Malformed JSON should be skipped
+    let input = r#"<｜tool▁calls▁begin｜>
+<｜tool▁call▁begin｜>function<｜tool▁sep｜>broken
+```json
+{invalid json}
+```<｜tool▁call▁end｜>
+<｜tool▁call▁begin｜>function<｜tool▁sep｜>valid
+```json
+{"key": "value"}
+```<｜tool▁call▁end｜>
+<｜tool▁calls▁end｜>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    // Only the valid tool call should be parsed
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "valid");
+}
+
+#[tokio::test]
+async fn test_multiple_tool_calls() {
+    let parser = DeepSeekParser::new();
+
+    let input = r#"<｜tool▁calls▁begin｜>
+<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather
+```json
+{"location": "Tokyo"}
+```<｜tool▁call▁end｜>
+<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather
+```json
+{"location": "Paris"}
+```<｜tool▁call▁end｜>
+<｜tool▁calls▁end｜><｜end▁of▁sentence｜>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 2);
+    assert_eq!(tools[0].function.name, "get_weather");
+    assert_eq!(tools[1].function.name, "get_weather");
+}

sgl-model-gateway/tests/tool_parser_edge_cases.rs

@@ -0,0 +1,351 @@
+//! Edge Cases and Error Handling Tests
+//!
+//! Tests for malformed input, edge cases, and error recovery
+
+use sgl_model_gateway::tool_parser::{
+    JsonParser, MistralParser, PythonicParser, QwenParser, ToolParser,
+};
+
+mod common;
+use common::create_test_tools;
+
+#[tokio::test]
+async fn test_empty_input() {
+    // Test that all parsers handle empty input correctly
+    let json_parser = JsonParser::new();
+    let (_normal_text, tools) = json_parser.parse_complete("").await.unwrap();
+    assert_eq!(
+        tools.len(),
+        0,
+        "JSON parser should return empty for empty input"
+    );
+
+    let mistral_parser = MistralParser::new();
+    let (_normal_text, tools) = mistral_parser.parse_complete("").await.unwrap();
+    assert_eq!(
+        tools.len(),
+        0,
+        "Mistral parser should return empty for empty input"
+    );
+
+    let qwen_parser = QwenParser::new();
+    let (_normal_text, tools) = qwen_parser.parse_complete("").await.unwrap();
+    assert_eq!(
+        tools.len(),
+        0,
+        "Qwen parser should return empty for empty input"
+    );
+
+    let pythonic_parser = PythonicParser::new();
+    let (_normal_text, tools) = pythonic_parser.parse_complete("").await.unwrap();
+    assert_eq!(
+        tools.len(),
+        0,
+        "Pythonic parser should return empty for empty input"
+    );
+}
+
+#[tokio::test]
+async fn test_plain_text_no_tools() {
+    let plain_text = "This is just a regular response with no tool calls whatsoever.";
+
+    let json_parser = JsonParser::new();
+    assert_eq!(
+        json_parser
+            .parse_complete(plain_text)
+            .await
+            .unwrap()
+            .1
+            .len(),
+        0
+    );
+
+    let mistral_parser = MistralParser::new();
+    assert_eq!(
+        mistral_parser
+            .parse_complete(plain_text)
+            .await
+            .unwrap()
+            .1
+            .len(),
+        0
+    );
+
+    let qwen_parser = QwenParser::new();
+    assert_eq!(
+        qwen_parser
+            .parse_complete(plain_text)
+            .await
+            .unwrap()
+            .1
+            .len(),
+        0
+    );
+
+    let pythonic_parser = PythonicParser::new();
+    assert_eq!(
+        pythonic_parser
+            .parse_complete(plain_text)
+            .await
+            .unwrap()
+            .1
+            .len(),
+        0
+    );
+}
+
+#[tokio::test]
+async fn test_incomplete_json() {
+    let json_parser = JsonParser::new();
+
+    let incomplete_cases = vec![
+        r#"{"name": "test""#,                 // Missing closing brace
+        r#"{"name": "test", "arguments":"#,   // Incomplete arguments
+        r#"{"name": "test", "arguments": {"#, // Incomplete nested object
+    ];
+
+    for input in incomplete_cases {
+        let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
+        assert_eq!(
+            tools.len(),
+            0,
+            "Should not parse incomplete JSON: {}",
+            input
+        );
+    }
+
+    // This case might actually parse because [{"name": "test"}] is complete
+    // The trailing comma suggests more items but the first item is valid
+    let _result = json_parser
+        .parse_complete(r#"[{"name": "test"},"#)
+        .await
+        .unwrap();
+    // This could parse the first element or return empty - implementation dependent
+}
+
+#[tokio::test]
+async fn test_malformed_mistral() {
+    let parser = MistralParser::new();
+
+    let malformed_cases = vec![
+        "[TOOL_CALLS]",                // Missing array
+        "[TOOL_CALLS] {",              // Not an array
+        "[TOOL_CALLS] [",              // Incomplete array
+        "[TOOL_CALLS] [{]",            // Invalid JSON in array
+        "[TOOL_CALLS] [{\"name\": }]", // Invalid value
+    ];
+
+    for input in malformed_cases {
+        // Parser might return error or empty vec for malformed input
+        if let Ok((_normal_text, tools)) = parser.parse_complete(input).await {
+            assert_eq!(
+                tools.len(),
+                0,
+                "Should not parse malformed Mistral: {}",
+                input
+            );
+        }
+        // Error is also acceptable for malformed input
+    }
+}
+
+#[tokio::test]
+async fn test_missing_required_fields() {
+    let json_parser = JsonParser::new();
+
+    // Missing name field
+    let input = r#"{"arguments": {"x": 1}}"#;
+    let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0, "Should not parse without name field");
+
+    // Name is not a string
+    let input = r#"{"name": 123, "arguments": {}}"#;
+    let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0, "Should not parse with non-string name");
+}
+
+#[tokio::test]
+async fn test_very_long_strings() {
+    let json_parser = JsonParser::new();
+
+    let long_string = "x".repeat(10000);
+    let input = format!(
+        r#"{{"name": "test", "arguments": {{"data": "{}"}}}}"#,
+        long_string
+    );
+
+    let (_normal_text, tools) = json_parser.parse_complete(&input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "test");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["data"].as_str().unwrap().len(), 10000);
+}
+
+#[tokio::test]
+async fn test_unicode_edge_cases() {
+    let json_parser = JsonParser::new();
+
+    // Various Unicode characters including emojis, CJK, RTL text
+    let input = r#"{"name": "translate", "arguments": {"text": "Hello 世界 🌍 مرحبا עולם"}}"#;
+
+    let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["text"], "Hello 世界 🌍 مرحبا עולם");
+}
+
+#[tokio::test]
+async fn test_nested_brackets_in_strings() {
+    let mistral_parser = MistralParser::new();
+    let input = r#"[TOOL_CALLS] [{"name": "echo", "arguments": {"text": "Array: [1, 2, 3]"}}]"#;
+    let (_normal_text, tools) = mistral_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["text"], "Array: [1, 2, 3]");
+
+    let pythonic_parser = PythonicParser::new();
+    let input = r#"[echo(text="List: [a, b, c]")]"#;
+    let (_normal_text, tools) = pythonic_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["text"], "List: [a, b, c]");
+}
+
+#[tokio::test]
+async fn test_multiple_formats_in_text() {
+    let json_parser = JsonParser::new();
+    let input = r#"
+    Here's some text with [TOOL_CALLS] that shouldn't trigger.
+    {"name": "actual_tool", "arguments": {}}
+    And some more text with <tool_call> tags.
+    "#;
+
+    let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "actual_tool");
+}
+
+#[tokio::test]
+async fn test_escaped_characters() {
+    let json_parser = JsonParser::new();
+
+    let input = r#"{"name": "write", "arguments": {"content": "Line 1\nLine 2\r\nLine 3\tTabbed\\Backslash\"Quote"}}"#;
+
+    let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    let content = args["content"].as_str().unwrap();
+    assert!(content.contains('\n'));
+    assert!(content.contains('\t'));
+    assert!(content.contains('\\'));
+    assert!(content.contains('"'));
+}
+
+#[tokio::test]
+async fn test_numeric_edge_cases() {
+    let json_parser = JsonParser::new();
+
+    let input = r#"{
+        "name": "calculate",
+        "arguments": {
+            "int": 42,
+            "float": 123.456,
+            "scientific": 1.23e-4,
+            "negative": -999,
+            "zero": 0,
+            "large": 9007199254740991
+        }
+    }"#;
+
+    let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["int"], 42);
+    assert_eq!(args["float"], 123.456);
+    assert_eq!(args["scientific"], 0.000123);
+    assert_eq!(args["negative"], -999);
+    assert_eq!(args["zero"], 0);
+    assert_eq!(args["large"], 9007199254740991i64);
+}
+
+#[tokio::test]
+async fn test_null_and_boolean_values() {
+    let json_parser = JsonParser::new();
+
+    let input = r#"{
+        "name": "configure",
+        "arguments": {
+            "enabled": true,
+            "disabled": false,
+            "optional": null
+        }
+    }"#;
+
+    let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["enabled"], true);
+    assert_eq!(args["disabled"], false);
+    assert_eq!(args["optional"], serde_json::Value::Null);
+}
+
+#[tokio::test]
+async fn test_partial_token_at_buffer_boundary() {
+    let mut parser = QwenParser::new();
+
+    let tools = create_test_tools();
+
+    // Send exactly "<tool" which is a 5-character prefix of "<tool_call>\n"
+    let result = parser.parse_incremental("<tool", &tools).await.unwrap();
+    assert!(
+        result.calls.is_empty(),
+        "Should be incomplete for partial tag"
+    );
+
+    // Complete the token
+    let result = parser
+        .parse_incremental(
+            "_call>\n{\"name\": \"test\", \"arguments\": {}}\n</tool_call>",
+            &tools,
+        )
+        .await
+        .unwrap();
+
+    // Should successfully parse after completing
+    if !result.calls.is_empty() {
+        if let Some(name) = &result.calls[0].name {
+            assert_eq!(name, "test");
+        }
+    }
+}
+
+#[tokio::test]
+async fn test_exact_prefix_lengths() {
+    let mut parser = QwenParser::new();
+
+    let tools = create_test_tools();
+
+    let test_cases = vec![
+        ("<", 1),            // 1-char prefix
+        ("<t", 2),           // 2-char prefix
+        ("<tool", 5),        // 5-char prefix (the main bug case)
+        ("<tool_call", 10),  // 10-char prefix
+        ("<tool_call>", 11), // 11-char prefix (full start without \n)
+    ];
+
+    for (prefix, expected_len) in test_cases {
+        let result = parser.parse_incremental(prefix, &tools).await.unwrap();
+        assert!(
+            result.calls.is_empty(),
+            "Prefix '{}' (len {}) should be incomplete",
+            prefix,
+            expected_len
+        );
+        // Buffer is now internal to parser - can't assert on it
+    }
+}

sgl-model-gateway/tests/tool_parser_fallback.rs

@@ -0,0 +1,272 @@
+//! Tests for tool parser fallback behavior
+//!
+//! When tool call parsing fails, the original text should be preserved as normal text
+//! rather than being lost. This ensures graceful degradation.
+
+use sgl_model_gateway::tool_parser::{
+    DeepSeekParser, JsonParser, LlamaParser, MistralParser, QwenParser, ToolParser,
+};
+
+#[tokio::test]
+async fn test_json_parser_invalid_json_returns_as_normal_text() {
+    let parser = JsonParser::new();
+
+    // Malformed JSON should be returned as normal text (note: commas may be processed)
+    let input = r#"{"name": "test", "arguments": invalid json here}"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(
+        normal_text,
+        r#"{"name": "test", "arguments": invalid json here}"#
+    );
+
+    // Plain text with no JSON structure should be returned as normal text
+    let input = "This is just plain text that should not be parsed as a tool call";
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input);
+
+    // Text that looks like it might have JSON but doesn't should be returned as normal text
+    let input = "The user said: {something} but it's not valid JSON";
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input);
+}
+
+#[tokio::test]
+async fn test_qwen_parser_invalid_format_returns_as_normal_text() {
+    let parser = QwenParser::new();
+
+    // Missing closing tag
+    let input = r#"<tool_call>
+{"name": "test", "arguments": {}}
+This text is missing the closing tag"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input); // Should preserve original text when no valid tools found
+
+    // Malformed JSON inside valid tags
+    let input = r#"<tool_call>
+{"name": "test", "arguments": invalid}
+</tool_call>"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    // When JSON parsing fails but tags are present, it should preserve the original text
+    assert_eq!(normal_text, input);
+
+    // Plain text without any tool markers
+    let input = "This is a regular response without any tool calls.";
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input); // Should return original text when no markers found
+}
+
+#[tokio::test]
+async fn test_llama_parser_invalid_format_returns_as_normal_text() {
+    let parser = LlamaParser::new();
+
+    // Invalid JSON after python_tag
+    let input = r#"<|python_tag|>{"name": "test", "arguments": invalid}"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input); // Should preserve original text when parsing fails
+
+    // Plain text without markers or JSON
+    let input = "Just explaining something without any function calls.";
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input); // Should return original text
+
+    // Text with python_tag but completely invalid content
+    let input = r#"Here's my response <|python_tag|>not even close to JSON"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input); // Should preserve everything when parsing fails
+}
+
+#[tokio::test]
+async fn test_mistral_parser_invalid_format_returns_as_normal_text() {
+    let parser = MistralParser::new();
+
+    // Missing closing bracket
+    let input = r#"[TOOL_CALLS] [{"name": "test", "arguments": {}"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input); // Should preserve original text when parsing fails
+
+    // Invalid JSON in tool calls section
+    let input = r#"[TOOL_CALLS] [{"name": invalid json}]"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input); // Should preserve original text when parsing fails
+
+    // Plain text
+    let input = "No tool calls here, just regular text.";
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input); // Should return original text
+}
+
+#[tokio::test]
+async fn test_deepseek_parser_invalid_format_returns_as_normal_text() {
+    let parser = DeepSeekParser::new();
+
+    // Invalid JSON in tool call
+    let input = r#"Some text<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>test
+```json
+{"name": "test", "arguments": malformed}
+```<｜tool▁call▁end｜><｜tool▁calls▁end｜>"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input); // Should preserve original text when parsing fails
+
+    // Missing function marker
+    let input = r#"<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>notfunction<｜tool▁sep｜>test
+```json
+{"x": 1}
+```<｜tool▁call▁end｜><｜tool▁calls▁end｜>"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input); // Should return original text when parsing fails
+
+    // No tool markers at all
+    let input = "Regular response without any special markers.";
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input); // Should return original text
+}
+
+#[tokio::test]
+async fn test_mixed_valid_and_invalid_content() {
+    let parser = QwenParser::new();
+
+    // Text with one valid tool call and one invalid
+    let input = r#"Let me help you with that.
+<tool_call>
+{"name": "valid_tool", "arguments": {"x": 1}}
+</tool_call>
+And here's another one:
+<tool_call>
+{"name": "invalid_tool", "arguments": malformed}
+</tool_call>
+That's all!"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1); // Should extract the valid tool
+    assert_eq!(tools[0].function.name, "valid_tool");
+    // Normal text should contain text before the first tool call
+    assert_eq!(normal_text, "Let me help you with that.\n");
+}
+
+#[tokio::test]
+async fn test_partial_tool_markers() {
+    // Test cases where tool markers are incomplete or cut off
+
+    let parser = QwenParser::new();
+    let input = "<tool_call>\nThis looks like it might be a tool call but it's not";
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input);
+
+    let parser = MistralParser::new();
+    let input = "[TOOL_CALLS] But then nothing follows...";
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input);
+
+    let parser = LlamaParser::new();
+    let input = "Starting a response <|python_tag|> but no JSON";
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input);
+}
+
+#[tokio::test]
+async fn test_escaped_json_like_content() {
+    // Test that JSON-like content in regular text doesn't get parsed as tools
+
+    let parser = JsonParser::new();
+    let input = r#"The user typed: {"name": "example"} but this is just quoted text"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    // JsonParser should extract the valid JSON and return normal text
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "example");
+    assert_eq!(normal_text, "The user typed:  but this is just quoted text");
+
+    let parser = QwenParser::new();
+    let input = r#"The syntax is: <tool_call>
+{"name": "example"}
+</tool_call> - that's how you format it"#;
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    // This actually contains valid tool call syntax, so it should parse
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "example");
+}
+
+#[tokio::test]
+async fn test_unicode_and_special_chars_in_failed_parsing() {
+    let parser = QwenParser::new();
+
+    // Unicode in malformed tool calls
+    let input = r#"<tool_call>
+{"name": "测试", "arguments": 🚀 invalid}
+</tool_call>"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    // Should handle Unicode properly in the fallback text - malformed content should be preserved
+    assert_eq!(normal_text, input);
+
+    // Special characters that might confuse parsers
+    let input = r#"Response: <tool_call>{"name": "test\n\t", "arguments": {"]}"}</tool_call>"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    // This might or might not parse depending on JSON handling of escape sequences
+    if tools.is_empty() {
+        assert!(!normal_text.is_empty() || normal_text == input);
+    }
+}
+
+#[tokio::test]
+async fn test_very_long_invalid_input() {
+    let parser = JsonParser::new();
+
+    // Generate a very long string that looks like it might be JSON but isn't
+    let mut input = String::from("{\"name\": \"test\", \"arguments\": {");
+    for i in 0..1000 {
+        input.push_str(&format!("\"field{}\": \"value{}\", ", i, i));
+    }
+    input.push_str("\"final\": incomplete"); // Don't close the JSON properly
+
+    let (normal_text, tools) = parser.parse_complete(&input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input); // Invalid JSON should be returned as normal text
+}
+
+#[tokio::test]
+async fn test_almost_valid_tool_calls() {
+    // Test tool calls that are almost valid but have small issues
+
+    let parser = JsonParser::new();
+
+    // Missing closing quote should be returned as normal text
+    let input = r#"{"name": "test", "arguments": {"key": "value}}"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(
+        normal_text,
+        r#"{"name": "test", "arguments": {"key": "value}}"#
+    );
+
+    // Extra comma
+    let input = r#"{"name": "test", "arguments": {},}"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    // Some JSON parsers might accept trailing commas
+    if tools.is_empty() {
+        assert_eq!(normal_text, r#"{"name": "test", "arguments": {},}"#);
+    }
+
+    // Wrong quote types
+    let input = r#"{'name': 'test', 'arguments': {}}"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0); // Standard JSON requires double quotes
+    assert_eq!(normal_text, r#"{'name': 'test', 'arguments': {}}"#);
+}

sgl-model-gateway/tests/tool_parser_glm4_moe.rs

@@ -0,0 +1,169 @@
+//! GLM-4 MoE Parser Integration Tests
+
+use sgl_model_gateway::tool_parser::{Glm4MoeParser, ToolParser};
+
+mod common;
+use common::create_test_tools;
+
+#[tokio::test]
+async fn test_glm4_complete_parsing() {
+    let parser = Glm4MoeParser::new();
+
+    let input = r#"Let me search for that.
+<tool_call>get_weather
+<arg_key>city</arg_key>
+<arg_value>Beijing</arg_value>
+<arg_key>date</arg_key>
+<arg_value>2024-12-25</arg_value>
+</tool_call>
+The weather will be..."#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "Let me search for that.\n");
+    assert_eq!(tools[0].function.name, "get_weather");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["city"], "Beijing");
+    assert_eq!(args["date"], "2024-12-25");
+}
+
+#[tokio::test]
+async fn test_glm4_multiple_tools() {
+    let parser = Glm4MoeParser::new();
+
+    let input = r#"<tool_call>search
+<arg_key>query</arg_key>
+<arg_value>rust tutorials</arg_value>
+</tool_call>
+<tool_call>translate
+<arg_key>text</arg_key>
+<arg_value>Hello World</arg_value>
+<arg_key>target_lang</arg_key>
+<arg_value>zh</arg_value>
+</tool_call>"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 2);
+    assert_eq!(normal_text, "");
+    assert_eq!(tools[0].function.name, "search");
+    assert_eq!(tools[1].function.name, "translate");
+}
+
+#[tokio::test]
+async fn test_glm4_type_conversion() {
+    let parser = Glm4MoeParser::new();
+
+    let input = r#"<tool_call>process
+<arg_key>count</arg_key>
+<arg_value>42</arg_value>
+<arg_key>rate</arg_key>
+<arg_value>1.5</arg_value>
+<arg_key>enabled</arg_key>
+<arg_value>true</arg_value>
+<arg_key>data</arg_key>
+<arg_value>null</arg_value>
+<arg_key>text</arg_key>
+<arg_value>string value</arg_value>
+</tool_call>"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["count"], 42);
+    assert_eq!(args["rate"], 1.5);
+    assert_eq!(args["enabled"], true);
+    assert_eq!(args["data"], serde_json::Value::Null);
+    assert_eq!(args["text"], "string value");
+}
+
+#[tokio::test]
+async fn test_glm4_streaming() {
+    let mut parser = Glm4MoeParser::new();
+
+    let tools = create_test_tools();
+
+    // Simulate streaming chunks
+    let chunks = vec![
+        "<tool_call>",
+        "get_weather\n",
+        "<arg_key>city</arg_key>\n",
+        "<arg_value>Shanghai</arg_value>\n",
+        "<arg_key>units</arg_key>\n",
+        "<arg_value>celsius</arg_value>\n",
+        "</tool_call>",
+    ];
+
+    let mut found_name = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+
+        for call in result.calls {
+            if let Some(name) = call.name {
+                assert_eq!(name, "get_weather");
+                found_name = true;
+            }
+        }
+    }
+
+    assert!(found_name, "Should have found tool name during streaming");
+}
+
+#[test]
+fn test_glm4_format_detection() {
+    let parser = Glm4MoeParser::new();
+
+    // Should detect GLM-4 format
+    assert!(parser.has_tool_markers("<tool_call>"));
+    assert!(parser.has_tool_markers("text with <tool_call> marker"));
+
+    // Should not detect other formats
+    assert!(!parser.has_tool_markers("[TOOL_CALLS]"));
+    assert!(!parser.has_tool_markers("<｜tool▁calls▁begin｜>"));
+    assert!(!parser.has_tool_markers("plain text"));
+}
+
+#[tokio::test]
+async fn test_python_literals() {
+    let parser = Glm4MoeParser::new();
+
+    let input = r#"<tool_call>test_func
+<arg_key>bool_true</arg_key>
+<arg_value>True</arg_value>
+<arg_key>bool_false</arg_key>
+<arg_value>False</arg_value>
+<arg_key>none_val</arg_key>
+<arg_value>None</arg_value>
+</tool_call>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "test_func");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["bool_true"], true);
+    assert_eq!(args["bool_false"], false);
+    assert_eq!(args["none_val"], serde_json::Value::Null);
+}
+
+#[tokio::test]
+async fn test_glm4_nested_json_in_arg_values() {
+    let parser = Glm4MoeParser::new();
+
+    let input = r#"<tool_call>process
+<arg_key>data</arg_key>
+<arg_value>{"nested": {"key": "value"}}</arg_value>
+<arg_key>list</arg_key>
+<arg_value>[1, 2, 3]</arg_value>
+</tool_call>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert!(args["data"].is_object());
+    assert!(args["list"].is_array());
+}

sgl-model-gateway/tests/tool_parser_json.rs

@@ -0,0 +1,717 @@
+//! JSON Parser Integration Tests
+//!
+//! Tests for the JSON parser which handles OpenAI, Claude, and generic JSON formats
+
+use serde_json::json;
+use sgl_model_gateway::tool_parser::{JsonParser, ToolParser};
+
+mod common;
+use common::{create_test_tools, streaming_helpers::*};
+
+#[tokio::test]
+async fn test_simple_json_tool_call() {
+    let parser = JsonParser::new();
+    let input = r#"{"name": "get_weather", "arguments": {"location": "San Francisco"}}"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "");
+    assert_eq!(tools[0].function.name, "get_weather");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["location"], "San Francisco");
+}
+
+#[tokio::test]
+async fn test_json_array_of_tools() {
+    let parser = JsonParser::new();
+    let input = r#"Hello, here are the results: [
+        {"name": "get_weather", "arguments": {"location": "SF"}},
+        {"name": "search", "arguments": {"query": "news"}}
+    ]"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 2);
+    assert_eq!(normal_text, "Hello, here are the results: ");
+    assert_eq!(tools[0].function.name, "get_weather");
+    assert_eq!(tools[1].function.name, "search");
+}
+
+#[tokio::test]
+async fn test_json_with_parameters_key() {
+    let parser = JsonParser::new();
+    let input = r#"{"name": "calculate", "parameters": {"x": 10, "y": 20}}"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "");
+    assert_eq!(tools[0].function.name, "calculate");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["x"], 10);
+    assert_eq!(args["y"], 20);
+}
+
+#[tokio::test]
+async fn test_json_extraction_from_text() {
+    let parser = JsonParser::new();
+    let input = r#"I'll help you with that. {"name": "search", "arguments": {"query": "rust"}} Let me search for that."#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(
+        normal_text,
+        "I'll help you with that.  Let me search for that."
+    );
+    assert_eq!(tools[0].function.name, "search");
+}
+
+#[tokio::test]
+async fn test_json_with_nested_objects() {
+    let parser = JsonParser::new();
+    let input = r#"{
+        "name": "update_config",
+        "arguments": {
+            "settings": {
+                "theme": "dark",
+                "language": "en",
+                "notifications": {
+                    "email": true,
+                    "push": false
+                }
+            }
+        }
+    }"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "");
+    assert_eq!(tools[0].function.name, "update_config");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["settings"]["theme"], "dark");
+    assert_eq!(args["settings"]["notifications"]["email"], true);
+}
+
+#[tokio::test]
+async fn test_json_with_special_characters() {
+    let parser = JsonParser::new();
+    let input = r#"{"name": "echo", "arguments": {"text": "Line 1\nLine 2\tTabbed", "path": "C:\\Users\\test"}}"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["text"], "Line 1\nLine 2\tTabbed");
+    assert_eq!(args["path"], "C:\\Users\\test");
+}
+
+#[tokio::test]
+async fn test_json_with_unicode() {
+    let parser = JsonParser::new();
+    let input = r#"{"name": "translate", "arguments": {"text": "Hello 世界 🌍", "emoji": "😊"}}"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["text"], "Hello 世界 🌍");
+    assert_eq!(args["emoji"], "😊");
+}
+
+#[tokio::test]
+async fn test_json_empty_arguments() {
+    let parser = JsonParser::new();
+    let input = r#"{"name": "ping", "arguments": {}}"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "");
+    assert_eq!(tools[0].function.name, "ping");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args, json!({}));
+}
+
+#[tokio::test]
+async fn test_json_invalid_format() {
+    let parser = JsonParser::new();
+
+    // Missing closing brace
+    let input = r#"{"name": "test", "arguments": {"key": "value""#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(
+        normal_text,
+        "{\"name\": \"test\", \"arguments\": {\"key\": \"value\""
+    );
+
+    // Not JSON at all
+    let input = "This is just plain text";
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+}
+
+#[tokio::test]
+async fn test_json_format_detection() {
+    let parser = JsonParser::new();
+
+    assert!(parser.has_tool_markers(r#"{"name": "test", "arguments": {}}"#));
+    assert!(parser.has_tool_markers(r#"[{"name": "test"}]"#));
+    assert!(!parser.has_tool_markers("plain text"));
+}
+
+// Streaming tests for JSON array format
+#[tokio::test]
+async fn test_json_array_streaming_required_mode() {
+    use sgl_model_gateway::protocols::common::Tool;
+
+    // Test that simulates the exact streaming pattern from required mode
+    let mut parser = JsonParser::new();
+
+    // Define test tools
+    let tools = vec![Tool {
+        tool_type: "function".to_string(),
+        function: sgl_model_gateway::protocols::common::Function {
+            name: "get_weather".to_string(),
+            description: Some("Get weather".to_string()),
+            parameters: serde_json::json!({}),
+            strict: None,
+        },
+    }];
+
+    // Simulate the EXACT chunks from the debug log
+    let chunks = vec![
+        "[{",
+        " \"",
+        "name",
+        "\":",
+        " \"",
+        "get",
+        "_weather",
+        "\",",
+        " \"",
+        "parameters",
+        "\":",
+        " {",
+        " \"",
+        "city",
+        "\":",
+        " \"",
+        "Paris",
+        "\"",
+        " }",
+        " }]",
+    ];
+
+    let mut all_results = Vec::new();
+    let mut all_normal_text = String::new();
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        all_results.extend(result.calls);
+        all_normal_text.push_str(&result.normal_text);
+    }
+
+    // We should have gotten tool call chunks
+    assert!(
+        !all_results.is_empty(),
+        "Should have emitted tool call chunks"
+    );
+
+    // Should not have emitted any normal text (including the closing ])
+    assert_eq!(
+        all_normal_text, "",
+        "Should not emit normal text for JSON array format"
+    );
+
+    // Check that we got the function name
+    let has_name = all_results
+        .iter()
+        .any(|item| item.name.as_ref().is_some_and(|n| n == "get_weather"));
+    assert!(has_name, "Should have emitted function name");
+
+    // Check that we got the parameters
+    let has_params = all_results.iter().any(|item| !item.parameters.is_empty());
+    assert!(has_params, "Should have emitted parameters");
+}
+
+#[tokio::test]
+async fn test_json_array_multiple_tools_streaming() {
+    use sgl_model_gateway::protocols::common::Tool;
+
+    // Test with multiple tools in array
+    let mut parser = JsonParser::new();
+
+    let tools = vec![
+        Tool {
+            tool_type: "function".to_string(),
+            function: sgl_model_gateway::protocols::common::Function {
+                name: "get_weather".to_string(),
+                description: Some("Get weather".to_string()),
+                parameters: serde_json::json!({}),
+                strict: None,
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: sgl_model_gateway::protocols::common::Function {
+                name: "get_news".to_string(),
+                description: Some("Get news".to_string()),
+                parameters: serde_json::json!({}),
+                strict: None,
+            },
+        },
+    ];
+
+    // Split into smaller, more realistic chunks
+    let chunks = vec![
+        "[{",
+        "\"name\":",
+        "\"get_weather\"",
+        ",\"parameters\":",
+        "{\"city\":",
+        "\"SF\"}",
+        "}",
+        ",",
+        "{\"name\":",
+        "\"get_news\"",
+        ",\"parameters\":",
+        "{\"topic\":",
+        "\"tech\"}",
+        "}]",
+    ];
+
+    let mut all_results = Vec::new();
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        all_results.extend(result.calls);
+    }
+
+    // Should have gotten tool calls for both functions
+    let has_weather = all_results
+        .iter()
+        .any(|item| item.name.as_ref().is_some_and(|n| n == "get_weather"));
+    let has_news = all_results
+        .iter()
+        .any(|item| item.name.as_ref().is_some_and(|n| n == "get_news"));
+
+    assert!(has_weather, "Should have get_weather tool call");
+    assert!(has_news, "Should have get_news tool call");
+}
+
+#[tokio::test]
+async fn test_json_array_closing_bracket_separate_chunk() {
+    use sgl_model_gateway::protocols::common::Tool;
+
+    // Test case where the closing ] comes as a separate chunk
+    let mut parser = JsonParser::new();
+
+    let tools = vec![Tool {
+        tool_type: "function".to_string(),
+        function: sgl_model_gateway::protocols::common::Function {
+            name: "get_weather".to_string(),
+            description: Some("Get weather".to_string()),
+            parameters: json!({}),
+            strict: None,
+        },
+    }];
+
+    // Closing ] as separate chunk, followed by normal text
+    let chunks = vec![
+        "[{",
+        "\"",
+        "name",
+        "\":",
+        "\"",
+        "get",
+        "_weather",
+        "\",",
+        "\"",
+        "parameters",
+        "\":",
+        "{",
+        "\"",
+        "city",
+        "\":",
+        "\"",
+        "Paris",
+        "\"",
+        "}",
+        "}",
+        "]",
+        " Here's",
+        " the",
+        " weather",
+        " info",
+    ];
+
+    let mut all_normal_text = String::new();
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        all_normal_text.push_str(&result.normal_text);
+    }
+
+    // Should emit only the third chunk as normal text, NOT the ]
+    assert_eq!(
+        all_normal_text, " Here's the weather info",
+        "Should emit only normal text without ], got: '{}'",
+        all_normal_text
+    );
+}
+
+#[tokio::test]
+async fn test_json_single_object_with_trailing_text() {
+    use sgl_model_gateway::protocols::common::Tool;
+
+    // Test single object format (no array) with trailing text
+    let mut parser = JsonParser::new();
+
+    let tools = vec![Tool {
+        tool_type: "function".to_string(),
+        function: sgl_model_gateway::protocols::common::Function {
+            name: "get_weather".to_string(),
+            description: Some("Get weather".to_string()),
+            parameters: serde_json::json!({}),
+            strict: None,
+        },
+    }];
+
+    let chunks = vec![
+        "{",
+        "\"",
+        "name",
+        "\":",
+        "\"",
+        "get_weather",
+        "\",",
+        "\"",
+        "parameters",
+        "\":",
+        "{",
+        "\"city",
+        "\":",
+        "\"Paris",
+        "\"}",
+        "}",
+        " Here's",
+        " the",
+        " weather",
+    ];
+
+    let mut all_normal_text = String::new();
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        all_normal_text.push_str(&result.normal_text);
+    }
+
+    // Should emit the trailing text as normal_text (no ] to strip for single object)
+    assert_eq!(
+        all_normal_text, " Here's the weather",
+        "Should emit normal text for single object format, got: '{}'",
+        all_normal_text
+    );
+}
+
+#[tokio::test]
+async fn test_json_single_object_with_bracket_in_text() {
+    use sgl_model_gateway::protocols::common::Tool;
+
+    // Test that ] in normal text is NOT stripped for single object format
+    let mut parser = JsonParser::new();
+
+    let tools = vec![Tool {
+        tool_type: "function".to_string(),
+        function: sgl_model_gateway::protocols::common::Function {
+            name: "get_weather".to_string(),
+            description: Some("Get weather".to_string()),
+            parameters: serde_json::json!({}),
+            strict: None,
+        },
+    }];
+
+    let chunks = vec![
+        "{",
+        "\"name",
+        "\":",
+        "\"get_weather",
+        "\",",
+        "\"parameters",
+        "\":",
+        "{",
+        "\"city",
+        "\":",
+        "\"Paris",
+        "\"}",
+        "}",
+        "]",
+        " Here's",
+        " the",
+        " weather",
+    ];
+
+    let mut all_normal_text = String::new();
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        all_normal_text.push_str(&result.normal_text);
+    }
+
+    // For single object format, ] should NOT be stripped (it's part of normal text)
+    assert_eq!(
+        all_normal_text, "] Here's the weather",
+        "Should preserve ] in normal text for single object format, got: '{}'",
+        all_normal_text
+    );
+}
+
+#[tokio::test]
+async fn test_json_array_bracket_in_text_after_tools() {
+    use sgl_model_gateway::protocols::common::Tool;
+
+    // Test that ] in normal text AFTER array tools is preserved
+    let mut parser = JsonParser::new();
+
+    let tools = vec![Tool {
+        tool_type: "function".to_string(),
+        function: sgl_model_gateway::protocols::common::Function {
+            name: "get_weather".to_string(),
+            description: Some("Get weather".to_string()),
+            parameters: serde_json::json!({}),
+            strict: None,
+        },
+    }];
+
+    let chunks = vec![
+        "[",
+        "{",
+        "\"name",
+        "\":",
+        "\"get_weather",
+        "\",",
+        "\"parameters",
+        "\":",
+        "{",
+        "\"city",
+        "\":",
+        "\"Paris",
+        "\"}",
+        "}",
+        "]",
+        " Array",
+        " notation:",
+        " arr",
+        "[",
+        "0",
+        "]",
+    ];
+
+    let mut all_normal_text = String::new();
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        all_normal_text.push_str(&result.normal_text);
+    }
+
+    // Should preserve ] in normal text after array tools complete
+    assert_eq!(
+        all_normal_text, " Array notation: arr[0]",
+        "Should preserve ] in normal text after array tools, got: '{}'",
+        all_normal_text
+    );
+}
+// =============================================================================
+// REALISTIC STREAMING TESTS
+// =============================================================================
+
+#[tokio::test]
+async fn test_json_bug_incomplete_tool_name_string() {
+    let tools = create_test_tools();
+    let mut parser = JsonParser::new();
+
+    // This exact sequence triggered the bug:
+    // Parser receives {"name": " and must NOT parse it as empty name
+    let chunks = vec![
+        r#"{"#,
+        r#"""#,
+        r#"name"#,
+        r#"""#,
+        r#":"#,
+        r#" "#,
+        r#"""#, // ← Critical moment: parser has {"name": "
+        // At this point, partial_json should NOT allow incomplete strings
+        // when current_tool_name_sent=false
+        r#"search"#, // Use valid tool name from create_test_tools()
+        r#"""#,
+        r#", "#,
+        r#"""#,
+        r#"arguments"#,
+        r#"""#,
+        r#": {"#,
+        r#"""#,
+        r#"query"#,
+        r#"""#,
+        r#": "#,
+        r#"""#,
+        r#"rust programming"#,
+        r#"""#,
+        r#"}}"#,
+    ];
+
+    let mut got_tool_name = false;
+    let mut saw_empty_name = false;
+
+    for chunk in chunks.iter() {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+
+        for call in result.calls {
+            if let Some(name) = &call.name {
+                if name.is_empty() {
+                    saw_empty_name = true;
+                }
+                if name == "search" {
+                    got_tool_name = true;
+                }
+            }
+        }
+    }
+
+    assert!(
+        !saw_empty_name,
+        "Parser should NEVER return empty tool name"
+    );
+    assert!(got_tool_name, "Should have parsed tool name correctly");
+}
+
+#[tokio::test]
+async fn test_json_realistic_chunks_simple_tool() {
+    let tools = create_test_tools();
+    let mut parser = JsonParser::new();
+
+    let input = r#"{"name": "get_weather", "arguments": {"city": "Paris"}}"#;
+    let chunks = create_realistic_chunks(input);
+
+    assert!(chunks.len() > 10, "Should have many small chunks");
+
+    let mut got_tool_name = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
+        for call in result.calls {
+            if let Some(name) = call.name {
+                assert_eq!(name, "get_weather");
+                got_tool_name = true;
+            }
+        }
+    }
+
+    assert!(got_tool_name, "Should have parsed tool name");
+}
+
+#[tokio::test]
+async fn test_json_strategic_chunks_with_quotes() {
+    let tools = create_test_tools();
+    let mut parser = JsonParser::new();
+
+    let input = r#"{"name": "search", "arguments": {"query": "rust programming"}}"#;
+    let chunks = create_strategic_chunks(input);
+
+    // Strategic chunks break after quotes and colons
+    assert!(chunks.iter().any(|c| c.ends_with('"')));
+
+    let mut got_tool_name = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
+        for call in result.calls {
+            if call.name.is_some() {
+                got_tool_name = true;
+            }
+        }
+    }
+
+    assert!(got_tool_name, "Should have parsed tool name");
+}
+
+#[tokio::test]
+async fn test_json_incremental_arguments_streaming() {
+    let tools = create_test_tools();
+    let mut parser = JsonParser::new();
+
+    let input = r#"{"name": "search", "arguments": {"query": "test", "limit": 10}}"#;
+    let chunks = create_realistic_chunks(input);
+
+    let mut tool_name_sent = false;
+    let mut got_arguments = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
+        for call in result.calls {
+            if call.name.is_some() {
+                tool_name_sent = true;
+            }
+            if tool_name_sent && !call.parameters.is_empty() {
+                got_arguments = true;
+            }
+        }
+    }
+
+    assert!(tool_name_sent, "Should have sent tool name");
+    assert!(got_arguments, "Should have sent arguments");
+}
+
+#[tokio::test]
+async fn test_json_very_long_url_in_arguments() {
+    let tools = create_test_tools();
+    let mut parser = JsonParser::new();
+
+    // Simulate long URL arriving in many chunks
+    let long_url = "https://example.com/very/long/path/".to_string() + &"segment/".repeat(50);
+    let input = format!(
+        r#"{{"name": "search", "arguments": {{"query": "{}"}}}}"#,
+        long_url
+    );
+    let chunks = create_realistic_chunks(&input);
+
+    assert!(chunks.len() > 100, "Long URL should create many chunks");
+
+    let mut got_tool_name = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
+        for call in result.calls {
+            if call.name.is_some() {
+                got_tool_name = true;
+            }
+        }
+    }
+
+    assert!(got_tool_name, "Should have parsed tool name");
+}
+
+#[tokio::test]
+async fn test_json_unicode() {
+    let tools = create_test_tools();
+    let mut parser = JsonParser::new();
+
+    let input = r#"{"name": "search", "arguments": {"query": "Hello 世界 🌍"}}"#;
+    let chunks = create_realistic_chunks(input);
+
+    let mut got_tool_name = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
+        for call in result.calls {
+            if call.name.is_some() {
+                got_tool_name = true;
+            }
+        }
+    }
+
+    assert!(got_tool_name, "Should have parsed with unicode");
+}

sgl-model-gateway/tests/tool_parser_kimik2.rs

@@ -0,0 +1,157 @@
+//! Kimi K2 Parser Integration Tests
+
+use sgl_model_gateway::tool_parser::{KimiK2Parser, ToolParser};
+
+mod common;
+use common::create_test_tools;
+
+#[tokio::test]
+async fn test_kimik2_complete_parsing() {
+    let parser = KimiK2Parser::new();
+
+    let input = r#"Let me help you with that.
+<|tool_calls_section_begin|>
+<|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{"location": "Tokyo", "units": "celsius"}<|tool_call_end|>
+<|tool_calls_section_end|>
+The weather in Tokyo is..."#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "Let me help you with that.\n");
+    assert_eq!(tools[0].function.name, "get_weather");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["location"], "Tokyo");
+    assert_eq!(args["units"], "celsius");
+}
+
+#[tokio::test]
+async fn test_kimik2_multiple_tools() {
+    let parser = KimiK2Parser::new();
+
+    let input = r#"<|tool_calls_section_begin|>
+<|tool_call_begin|>functions.search:0<|tool_call_argument_begin|>{"query": "rust tutorials"}<|tool_call_end|>
+<|tool_call_begin|>functions.translate:1<|tool_call_argument_begin|>{"text": "Hello", "to": "ja"}<|tool_call_end|>
+<|tool_calls_section_end|>"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 2);
+    assert_eq!(normal_text, "");
+    assert_eq!(tools[0].function.name, "search");
+    assert_eq!(tools[1].function.name, "translate");
+}
+
+#[tokio::test]
+async fn test_kimik2_with_whitespace() {
+    let parser = KimiK2Parser::new();
+
+    let input = r#"<|tool_calls_section_begin|>
+<|tool_call_begin|> functions.test:0 <|tool_call_argument_begin|> {"key": "value", "num": 42} <|tool_call_end|>
+<|tool_calls_section_end|>"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "");
+    assert_eq!(tools[0].function.name, "test");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["key"], "value");
+    assert_eq!(args["num"], 42);
+}
+
+#[tokio::test]
+async fn test_kimik2_streaming() {
+    let tools = create_test_tools();
+
+    let mut parser = KimiK2Parser::new();
+
+    // Simulate streaming chunks
+    let chunks = vec![
+        "<|tool_calls_section_begin|>\n",
+        "<|tool_call_begin|>functions.",
+        "calculate:0",
+        "<|tool_call_argument_begin|>",
+        r#"{"x": 10, "#,
+        r#""y": 20}"#,
+        "<|tool_call_end|>\n",
+        "<|tool_calls_section_end|>",
+    ];
+
+    let mut found_name = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+
+        for call in result.calls {
+            if let Some(name) = call.name {
+                assert_eq!(name, "calculate");
+                found_name = true;
+            }
+        }
+    }
+
+    assert!(found_name, "Should have found tool name during streaming");
+}
+
+#[test]
+fn test_kimik2_format_detection() {
+    let parser = KimiK2Parser::new();
+
+    // Should detect Kimi K2 format
+    assert!(parser.has_tool_markers("<|tool_calls_section_begin|>"));
+    assert!(parser.has_tool_markers("text with <|tool_calls_section_begin|> marker"));
+
+    // Should not detect other formats
+    assert!(!parser.has_tool_markers("[TOOL_CALLS]"));
+    assert!(!parser.has_tool_markers("<tool_call>"));
+    assert!(!parser.has_tool_markers("plain text"));
+}
+
+#[tokio::test]
+async fn test_kimik2_sequential_indices() {
+    let parser = KimiK2Parser::new();
+
+    let input = r#"<|tool_calls_section_begin|>
+<|tool_call_begin|>functions.first:0<|tool_call_argument_begin|>{"param": "a"}<|tool_call_end|>
+<|tool_call_begin|>functions.second:1<|tool_call_argument_begin|>{"param": "b"}<|tool_call_end|>
+<|tool_call_begin|>functions.third:2<|tool_call_argument_begin|>{"param": "c"}<|tool_call_end|>
+<|tool_calls_section_end|>"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 3);
+    assert_eq!(normal_text, "");
+    assert_eq!(tools[0].function.name, "first");
+    assert_eq!(tools[1].function.name, "second");
+    assert_eq!(tools[2].function.name, "third");
+}
+
+#[tokio::test]
+async fn test_function_index_extraction() {
+    let parser = KimiK2Parser::new();
+
+    let input = r#"Text before tool calls.
+<|tool_calls_section_begin|>
+<|tool_call_begin|>functions.search:0<|tool_call_argument_begin|>{"query": "rust"}<|tool_call_end|>
+<|tool_call_begin|>functions.calc:1<|tool_call_argument_begin|>{"x": 10}<|tool_call_end|>
+<|tool_calls_section_end|>"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 2);
+    assert_eq!(normal_text, "Text before tool calls.\n");
+    assert_eq!(tools[0].function.name, "search");
+    assert_eq!(tools[1].function.name, "calc");
+    // TODO: Verify indices are preserved: 0 and 1
+}
+
+#[tokio::test]
+async fn test_namespace_extraction() {
+    let parser = KimiK2Parser::new();
+
+    let input = r#"<|tool_calls_section_begin|>
+<|tool_call_begin|>api.tools.search:0<|tool_call_argument_begin|>{"q": "test"}<|tool_call_end|>
+<|tool_calls_section_end|>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "api.tools.search"); // Includes full namespace
+}

sgl-model-gateway/tests/tool_parser_llama.rs

@@ -0,0 +1,455 @@
+//! Llama Parser Integration Tests
+//!
+//! Tests for the Llama parser which handles <|python_tag|> format and plain JSON
+
+use sgl_model_gateway::tool_parser::{LlamaParser, ToolParser};
+
+mod common;
+use common::{create_test_tools, streaming_helpers::*};
+
+#[tokio::test]
+async fn test_llama_python_tag_format() {
+    let parser = LlamaParser::new();
+    let input = r#"Here are some results: <|python_tag|>{"name": "search", "parameters": {"query": "weather"}}"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "search");
+    assert_eq!(normal_text, "Here are some results: ");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["query"], "weather");
+}
+
+#[tokio::test]
+async fn test_llama_with_semicolon_separation() {
+    let parser = LlamaParser::new();
+
+    let input = r#"<|python_tag|>{"name": "tool1", "parameters": {}};{"name": "tool2", "parameters": {"y": 2}}"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 2);
+    assert_eq!(tools[0].function.name, "tool1");
+    assert_eq!(tools[1].function.name, "tool2");
+    assert_eq!(normal_text, "");
+}
+
+#[tokio::test]
+async fn test_llama_no_tool_calls() {
+    let parser = LlamaParser::new();
+
+    let input = "This is just plain text with no tool calls";
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input);
+}
+
+#[tokio::test]
+async fn test_llama_plain_json_fallback() {
+    let parser = LlamaParser::new();
+    let input = r#"{"name": "calculate", "parameters": {"x": 5, "y": 10}}"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "calculate");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["x"], 5);
+    assert_eq!(args["y"], 10);
+}
+
+#[tokio::test]
+async fn test_llama_with_text_before() {
+    let parser = LlamaParser::new();
+    let input = r#"Let me help you with that. <|python_tag|>{"name": "get_time", "parameters": {"timezone": "UTC"}}"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "Let me help you with that. ");
+    assert_eq!(tools[0].function.name, "get_time");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["timezone"], "UTC");
+}
+
+#[tokio::test]
+async fn test_llama_with_nested_json() {
+    let parser = LlamaParser::new();
+    let input = r#"<|python_tag|>{
+        "name": "update_settings",
+        "parameters": {
+            "preferences": {
+                "theme": "dark",
+                "language": "en"
+            },
+            "notifications": true
+        }
+    }"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "update_settings");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["preferences"]["theme"], "dark");
+    assert_eq!(args["notifications"], true);
+}
+
+#[tokio::test]
+async fn test_llama_empty_arguments() {
+    let parser = LlamaParser::new();
+
+    // With python_tag
+    let input = r#"<|python_tag|>{"name": "ping", "parameters": {}}"#;
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "ping");
+
+    // Plain JSON
+    let input = r#"{"name": "ping", "parameters": {}}"#;
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "ping");
+}
+
+#[tokio::test]
+async fn test_llama_format_detection() {
+    let parser = LlamaParser::new();
+
+    assert!(parser.has_tool_markers(r#"<|python_tag|>{"name": "test"}"#));
+    assert!(parser.has_tool_markers(r#"{"name": "test", "parameters": {}}"#));
+    assert!(!parser.has_tool_markers("plain text"));
+}
+
+#[tokio::test]
+async fn test_llama_invalid_json_after_tag() {
+    let parser = LlamaParser::new();
+
+    let input = r#"<|python_tag|>{"name": invalid}"#;
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, "<|python_tag|>{\"name\": invalid}");
+}
+
+#[tokio::test]
+async fn test_llama_real_world_output() {
+    let parser = LlamaParser::new();
+
+    // Actual output from Llama 3.2 model - simplified for testing
+    let input = r#"I'll search for that information for you.
+
+<|python_tag|>{"name": "web_search", "parameters": {"query": "Llama 3.2 model capabilities", "num_results": 5, "search_type": "recent"}}"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "web_search");
+
+    let formatted_input = r#"<|python_tag|>{
+    "name": "get_current_time",
+    "parameters": {
+        "timezone": "America/New_York",
+        "format": "ISO8601"
+    }
+}"#;
+
+    let (_normal_text, tools2) = parser.parse_complete(formatted_input).await.unwrap();
+    assert_eq!(tools2.len(), 1);
+    assert_eq!(tools2[0].function.name, "get_current_time");
+}
+
+#[tokio::test]
+async fn test_single_json() {
+    let parser = LlamaParser::new();
+    let text = r#"{"name": "get_weather", "parameters": {"city": "Paris"}}"#;
+
+    let (_normal_text, tools) = parser.parse_complete(text).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "get_weather");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["city"], "Paris");
+}
+
+#[tokio::test]
+async fn test_multiple_json_with_separator() {
+    let parser = LlamaParser::new();
+    let text = r#"<|python_tag|>{"name": "get_weather", "parameters": {"city": "Paris"}};{"name": "get_tourist_attractions", "parameters": {"city": "Paris"}}"#;
+
+    let (_normal_text, tools) = parser.parse_complete(text).await.unwrap();
+    // Note: Current implementation may only parse the first one due to semicolon handling
+    assert!(!tools.is_empty());
+    assert_eq!(tools[0].function.name, "get_weather");
+}
+
+#[tokio::test]
+async fn test_json_with_trailing_text() {
+    let parser = LlamaParser::new();
+    // Valid JSON with trailing text - LlamaParser doesn't support this mixed format
+    let text = r#"{"name": "get_weather", "parameters": {}} Some follow-up text"#;
+
+    let (normal_text, tools) = parser.parse_complete(text).await.unwrap();
+    // LlamaParser expects pure JSON or <|python_tag|> format, not JSON with trailing text
+    // So this returns as normal text
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, text);
+}
+
+#[tokio::test]
+async fn test_invalid_then_valid_json() {
+    let parser = LlamaParser::new();
+    let text =
+        r#"{"name": "get_weather", "parameters": {{"name": "get_weather", "parameters": {}}"#;
+
+    let (_normal_text, tools) = parser.parse_complete(text).await.unwrap();
+    // Should parse at least one valid JSON
+    if !tools.is_empty() {
+        assert_eq!(tools[0].function.name, "get_weather");
+    }
+}
+
+#[tokio::test]
+async fn test_plain_text_only() {
+    let parser = LlamaParser::new();
+    let text = "This is just plain explanation text.";
+
+    let (_normal_text, tools) = parser.parse_complete(text).await.unwrap();
+    assert_eq!(tools.len(), 0);
+}
+
+#[tokio::test]
+async fn test_with_python_tag_prefix() {
+    let parser = LlamaParser::new();
+    let text = r#"Some intro. <|python_tag|>{"name": "get_weather", "parameters": {}}"#;
+
+    let (_normal_text, tools) = parser.parse_complete(text).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "get_weather");
+}
+
+// STREAMING TESTS
+
+#[tokio::test]
+async fn test_llama_streaming_simple() {
+    let tools = create_test_tools();
+
+    let mut parser = LlamaParser::new();
+
+    // Send complete JSON at once
+    let full_json = r#"<|python_tag|>{"name": "search", "parameters": {"query": "weather"}}"#;
+
+    let result = parser.parse_incremental(full_json, &tools).await.unwrap();
+
+    assert!(
+        !result.calls.is_empty(),
+        "Expected tool call for complete JSON input"
+    );
+    assert_eq!(result.calls[0].name.as_ref().unwrap(), "search");
+}
+
+#[tokio::test]
+async fn test_llama_streaming_partial() {
+    let tools = create_test_tools();
+
+    let mut parser = LlamaParser::new();
+
+    // Stream in chunks
+    let chunks = vec![
+        r#"<|python"#,
+        r#"_tag|>{"name": "#,
+        r#""calculate", "#,
+        r#""parameters": {"x": 10}"#,
+        r#"}"#,
+    ];
+
+    let mut got_complete = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        if !result.calls.is_empty() {
+            if let Some(name) = &result.calls[0].name {
+                assert_eq!(name, "calculate");
+                got_complete = true;
+            }
+        }
+    }
+
+    assert!(got_complete, "Should have completed parsing");
+}
+
+#[tokio::test]
+async fn test_llama_streaming_plain_json() {
+    let tools = create_test_tools();
+
+    let mut parser = LlamaParser::new();
+
+    // Stream plain JSON without python_tag
+    let chunks = vec![
+        r#"{"name": "#,
+        r#""search", "#,
+        r#""parameters": "#,
+        r#"{"query": "#,
+        r#""test"}}"#,
+    ];
+
+    let mut got_complete = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        if !result.calls.is_empty() {
+            if let Some(name) = &result.calls[0].name {
+                assert_eq!(name, "search");
+                got_complete = true;
+            }
+        }
+    }
+
+    assert!(got_complete, "Should have completed parsing");
+}
+
+#[tokio::test]
+async fn test_llama_streaming_with_text_before() {
+    let tools = create_test_tools();
+
+    let mut parser = LlamaParser::new();
+
+    let chunks = vec![
+        r#"Let me help you. "#,
+        r#"<|python_tag|>"#,
+        r#"{"name": "get_time","#,
+        r#" "parameters": {"#,
+        r#""timezone": "UTC"}}"#,
+    ];
+
+    let mut got_complete = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        if !result.calls.is_empty() {
+            if let Some(name) = &result.calls[0].name {
+                assert_eq!(name, "get_time");
+                got_complete = true;
+            }
+        }
+    }
+
+    assert!(got_complete, "Should have completed parsing");
+}
+
+#[tokio::test]
+async fn test_llama_streaming_multiple_tools() {
+    let tools = create_test_tools();
+
+    let mut parser = LlamaParser::new();
+
+    let text =
+        r#"<|python_tag|>{"name": "func1", "parameters": {}};{"name": "func2", "parameters": {}}"#;
+
+    let result = parser.parse_incremental(text, &tools).await.unwrap();
+
+    // Should get first tool complete
+    assert!(
+        !result.calls.is_empty(),
+        "Expected first tool to be complete"
+    );
+    if let Some(name) = &result.calls[0].name {
+        assert_eq!(name, "func1");
+    }
+
+    // Process remaining buffer to get second tool
+    let result2 = parser.parse_incremental("", &tools).await.unwrap();
+    if !result2.calls.is_empty() {
+        if let Some(name) = &result2.calls[0].name {
+            assert_eq!(name, "func2");
+        }
+    }
+}
+
+#[tokio::test]
+async fn test_llama_streaming_multiple_tools_chunked() {
+    let mut parser = LlamaParser::new();
+
+    let tools = create_test_tools();
+
+    // First chunk - incomplete first JSON
+    let chunk1 = r#"<|python_tag|>{"name": "get_weather", "parameters""#;
+    let result1 = parser.parse_incremental(chunk1, &tools).await.unwrap();
+    if !result1.calls.is_empty() {
+        if let Some(name) = &result1.calls[0].name {
+            assert_eq!(name, "get_weather");
+        }
+    }
+
+    // Second chunk - complete first JSON and separator
+    let chunk2 = r#": {"city": "Paris"}};{"name": "#;
+    let result2 = parser.parse_incremental(chunk2, &tools).await.unwrap();
+
+    // Should get parameters for first tool (name already sent in result1)
+    if !result2.calls.is_empty() {
+        let args: serde_json::Value = serde_json::from_str(&result2.calls[0].parameters).unwrap();
+        assert_eq!(args["city"], "Paris");
+    }
+
+    let chunk3 = r#""get_time", "parameters": {"timezone": "UTC"}}"#;
+    let result3 = parser.parse_incremental(chunk3, &tools).await.unwrap();
+    if !result3.calls.is_empty() {
+        if let Some(name) = &result3.calls[0].name {
+            assert_eq!(name, "get_time");
+        }
+    }
+}
+
+// =============================================================================
+// REALISTIC STREAMING TESTS
+// =============================================================================
+
+#[tokio::test]
+async fn test_llama_realistic_chunks_with_python_tag() {
+    let tools = create_test_tools();
+    let mut parser = LlamaParser::new();
+
+    let input = r#"<|python_tag|>{"name": "calculate", "parameters": {"x": 10, "y": 20}}"#;
+    let chunks = create_realistic_chunks(input);
+
+    assert!(chunks.len() > 15, "Should have many small chunks");
+
+    let mut got_tool_name = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
+        for call in result.calls {
+            if let Some(name) = call.name {
+                assert_eq!(name, "calculate");
+                got_tool_name = true;
+            }
+        }
+    }
+
+    assert!(got_tool_name, "Should have parsed tool name");
+}
+
+#[tokio::test]
+async fn test_llama_python_tag_arrives_in_parts() {
+    let tools = create_test_tools();
+    let mut parser = LlamaParser::new();
+
+    // Python tag itself arrives in small chunks
+    let chunks = vec![
+        "<|p", "yth", "on_", "tag", "|>{", r#"""#, "na", r#"me""#, ": ", r#"""#, "sea", "rch",
+        r#"""#, ", ", r#"""#, "par", "ame", "ter", "s", r#"""#, ": {", r#"""#, "q", r#"""#, ": ",
+        r#"""#, "tes", "t", r#"""#, "}}",
+    ];
+
+    let mut got_tool_name = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        for call in result.calls {
+            if let Some(name) = call.name {
+                assert_eq!(name, "search");
+                got_tool_name = true;
+            }
+        }
+    }
+
+    assert!(got_tool_name, "Should have parsed tool name");
+}

sgl-model-gateway/tests/tool_parser_minimax_m2.rs

@@ -0,0 +1,780 @@
+//! MiniMax M2 Parser Integration Tests
+
+use sgl_model_gateway::tool_parser::{MinimaxM2Parser, ToolParser};
+
+mod common;
+use common::create_test_tools;
+
+#[tokio::test]
+async fn test_minimax_complete_parsing() {
+    let parser = MinimaxM2Parser::new();
+
+    let input = r#"Let me search for that.
+<minimax:tool_call>
+<invoke name="get_weather">
+<parameter name="city">Beijing</parameter>
+<parameter name="date">2024-12-25</parameter>
+</invoke>
+</minimax:tool_call>
+The weather will be..."#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "Let me search for that.\n");
+    assert_eq!(tools[0].function.name, "get_weather");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["city"], "Beijing");
+    assert_eq!(args["date"], "2024-12-25");
+}
+
+#[tokio::test]
+async fn test_minimax_multiple_tools() {
+    let parser = MinimaxM2Parser::new();
+
+    let input = r#"<minimax:tool_call>
+<invoke name="search">
+<parameter name="query">rust tutorials</parameter>
+</invoke>
+</minimax:tool_call>
+<minimax:tool_call>
+<invoke name="translate">
+<parameter name="text">Hello World</parameter>
+<parameter name="target_lang">zh</parameter>
+</invoke>
+</minimax:tool_call>"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 2);
+    assert_eq!(normal_text, "");
+    assert_eq!(tools[0].function.name, "search");
+    assert_eq!(tools[1].function.name, "translate");
+}
+
+#[tokio::test]
+async fn test_minimax_type_conversion() {
+    let parser = MinimaxM2Parser::new();
+
+    let input = r#"<minimax:tool_call>
+<invoke name="process">
+<parameter name="count">42</parameter>
+<parameter name="rate">1.5</parameter>
+<parameter name="enabled">true</parameter>
+<parameter name="data">null</parameter>
+<parameter name="text">string value</parameter>
+</invoke>
+</minimax:tool_call>"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["count"], 42);
+    assert_eq!(args["rate"], 1.5);
+    assert_eq!(args["enabled"], true);
+    assert_eq!(args["data"], serde_json::Value::Null);
+    assert_eq!(args["text"], "string value");
+}
+
+#[tokio::test]
+async fn test_minimax_streaming_basic() {
+    let mut parser = MinimaxM2Parser::new();
+
+    let tools = create_test_tools();
+
+    // Simulate streaming chunks
+    let chunks = vec![
+        "<minimax:tool_call>",
+        r#"<invoke name="get_weather">"#,
+        r#"<parameter name="city">Shanghai</parameter>"#,
+        r#"<parameter name="units">celsius</parameter>"#,
+        "</invoke>",
+        "</minimax:tool_call>",
+    ];
+
+    let mut found_name = false;
+    let mut found_params = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+
+        for call in result.calls {
+            if let Some(name) = call.name {
+                assert_eq!(name, "get_weather");
+                found_name = true;
+            }
+            if !call.parameters.is_empty() {
+                found_params = true;
+            }
+        }
+    }
+
+    assert!(found_name, "Should have found tool name during streaming");
+    assert!(found_params, "Should have streamed parameters");
+}
+
+#[test]
+fn test_minimax_format_detection() {
+    let parser = MinimaxM2Parser::new();
+
+    // Should detect MiniMax format
+    assert!(parser.has_tool_markers("<minimax:tool_call>"));
+    assert!(parser.has_tool_markers("text with <minimax:tool_call> marker"));
+
+    // Should not detect other formats
+    assert!(!parser.has_tool_markers("<tool_call>")); // GLM4 format
+    assert!(!parser.has_tool_markers("[TOOL_CALLS]"));
+    assert!(!parser.has_tool_markers("<｜tool▁calls▁begin｜>"));
+    assert!(!parser.has_tool_markers("plain text"));
+}
+
+#[tokio::test]
+async fn test_minimax_python_literals() {
+    let parser = MinimaxM2Parser::new();
+
+    let input = r#"<minimax:tool_call>
+<invoke name="test_func">
+<parameter name="bool_true">True</parameter>
+<parameter name="bool_false">False</parameter>
+<parameter name="none_val">None</parameter>
+</invoke>
+</minimax:tool_call>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "test_func");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["bool_true"], true);
+    assert_eq!(args["bool_false"], false);
+    assert_eq!(args["none_val"], serde_json::Value::Null);
+}
+
+#[tokio::test]
+async fn test_minimax_nested_json_in_parameters() {
+    let parser = MinimaxM2Parser::new();
+
+    let input = r#"<minimax:tool_call>
+<invoke name="process">
+<parameter name="data">{"nested": {"key": "value"}}</parameter>
+<parameter name="list">[1, 2, 3]</parameter>
+</invoke>
+</minimax:tool_call>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    // JSON-like strings are kept as strings, not parsed as JSON
+    // This matches the behavior of other parsers like GLM4 MOE
+    assert!(args["data"].is_string());
+    assert_eq!(args["data"], r#"{"nested": {"key": "value"}}"#);
+    assert!(args["list"].is_string());
+    assert_eq!(args["list"], "[1, 2, 3]");
+}
+
+#[tokio::test]
+async fn test_minimax_xml_entities() {
+    let parser = MinimaxM2Parser::new();
+
+    let input = r#"<minimax:tool_call>
+<invoke name="process">
+<parameter name="html">&lt;div&gt;content&lt;/div&gt;</parameter>
+<parameter name="text">Quote: &quot;hello&quot;</parameter>
+<parameter name="code">if (a &amp;&amp; b) { }</parameter>
+</invoke>
+</minimax:tool_call>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["html"], "<div>content</div>");
+    assert_eq!(args["text"], "Quote: \"hello\"");
+    assert_eq!(args["code"], "if (a && b) { }");
+}
+
+#[tokio::test]
+async fn test_minimax_streaming_partial_tags() {
+    let mut parser = MinimaxM2Parser::new();
+    let tools = create_test_tools();
+
+    // Chunks split mid-tag
+    let chunks = vec![
+        "<minimax:tool_c",
+        "all><invoke na",
+        r#"me="get_weather"><param"#,
+        r#"eter name="city">Bei"#,
+        "jing</parameter></inv",
+        "oke></minimax:tool_call>",
+    ];
+
+    let mut found_name = false;
+    let mut buffer = String::new();
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+
+        buffer.push_str(&result.normal_text);
+
+        for call in result.calls {
+            if let Some(name) = call.name {
+                assert_eq!(name, "get_weather");
+                found_name = true;
+            }
+        }
+    }
+
+    assert!(
+        found_name,
+        "Should have parsed function name from partial chunks"
+    );
+    assert_eq!(buffer, "");
+}
+
+#[tokio::test]
+async fn test_minimax_streaming_incremental_json() {
+    let mut parser = MinimaxM2Parser::new();
+    let tools = create_test_tools();
+
+    let chunks = vec![
+        "<minimax:tool_call>",
+        r#"<invoke name="get_weather">"#,
+        r#"<parameter name="city">Paris</parameter>"#,
+        r#"<parameter name="units">metric</parameter>"#,
+        "</invoke></minimax:tool_call>",
+    ];
+
+    let mut json_fragments = Vec::new();
+    let mut found_function = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+
+        for call in result.calls {
+            if let Some(_name) = call.name {
+                found_function = true;
+            }
+            if !call.parameters.is_empty() {
+                json_fragments.push(call.parameters.clone());
+            }
+        }
+    }
+
+    assert!(found_function);
+
+    // Verify JSON was built incrementally
+    assert!(!json_fragments.is_empty());
+
+    // First fragment should start with opening brace
+    if let Some(first) = json_fragments.first() {
+        assert!(
+            first.starts_with('{'),
+            "First JSON fragment should start with '{{': {}",
+            first
+        );
+    }
+
+    // Last fragment should be closing brace
+    if let Some(last) = json_fragments.last() {
+        assert!(
+            last.contains('}'),
+            "Last JSON fragment should contain '}}': {}",
+            last
+        );
+    }
+}
+
+#[tokio::test]
+async fn test_minimax_multiple_tools_boundary() {
+    let mut parser = MinimaxM2Parser::new();
+    let tools = create_test_tools();
+
+    // Tool boundary at chunk boundary
+    let chunks = vec![
+        r#"<minimax:tool_call><invoke name="get_weather"><parameter name="city">Tokyo</parameter></invoke></minimax:tool_call>"#,
+        r#"<minimax:tool_call><invoke name="search"><parameter name="query">weather forecast</parameter></invoke></minimax:tool_call>"#,
+    ];
+
+    let mut tool_names = Vec::new();
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+
+        for call in result.calls {
+            if let Some(name) = call.name {
+                tool_names.push(name);
+            }
+        }
+    }
+
+    assert_eq!(tool_names.len(), 2);
+    assert_eq!(tool_names[0], "get_weather");
+    assert_eq!(tool_names[1], "search");
+}
+
+#[tokio::test]
+async fn test_minimax_invalid_function_name() {
+    let mut parser = MinimaxM2Parser::new();
+    let tools = create_test_tools();
+
+    let chunks = vec![
+        "<minimax:tool_call>",
+        r#"<invoke name="invalid_function">"#,
+        r#"<parameter name="param">value</parameter>"#,
+        "</invoke></minimax:tool_call>",
+    ];
+
+    let mut found_invalid = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+
+        // Invalid function should be skipped
+        for call in result.calls {
+            if let Some(name) = call.name {
+                if name == "invalid_function" {
+                    found_invalid = true;
+                }
+            }
+        }
+    }
+
+    assert!(!found_invalid, "Invalid function should not be parsed");
+}
+
+#[tokio::test]
+async fn test_minimax_empty_parameters() {
+    let parser = MinimaxM2Parser::new();
+
+    let input = r#"<minimax:tool_call>
+<invoke name="simple_func">
+</invoke>
+</minimax:tool_call>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "simple_func");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args, serde_json::json!({}));
+}
+
+#[tokio::test]
+async fn test_minimax_multiline_parameter_values() {
+    let parser = MinimaxM2Parser::new();
+
+    let input = r#"<minimax:tool_call>
+<invoke name="process">
+<parameter name="multiline">line1
+line2
+line3</parameter>
+<parameter name="unicode">你好世界 🌍</parameter>
+</invoke>
+</minimax:tool_call>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["multiline"], "line1\nline2\nline3");
+    assert_eq!(args["unicode"], "你好世界 🌍");
+}
+
+#[tokio::test]
+async fn test_minimax_nested_xml_like_content() {
+    let parser = MinimaxM2Parser::new();
+
+    let input = r#"<minimax:tool_call>
+<invoke name="process">
+<parameter name="template"><html><body>Hello</body></html></parameter>
+<parameter name="config">{"key": "<value>nested</value>"}</parameter>
+</invoke>
+</minimax:tool_call>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["template"], "<html><body>Hello</body></html>");
+
+    // The nested JSON with XML-like content
+    let config =
+        serde_json::from_str::<serde_json::Value>(args["config"].as_str().unwrap()).unwrap();
+    assert_eq!(config["key"], "<value>nested</value>");
+}
+
+#[tokio::test]
+async fn test_minimax_streaming_state_reset() {
+    let mut parser = MinimaxM2Parser::new();
+    let tools = create_test_tools();
+
+    // First tool
+    let chunks1 = vec![
+        r#"<minimax:tool_call><invoke name="get_weather">"#,
+        r#"<parameter name="city">London</parameter>"#,
+        "</invoke></minimax:tool_call>",
+    ];
+
+    for chunk in chunks1 {
+        parser.parse_incremental(chunk, &tools).await.unwrap();
+    }
+
+    // Second tool - state should be reset
+    let chunks2 = vec![
+        r#"<minimax:tool_call><invoke name="search">"#,
+        r#"<parameter name="query">rust</parameter>"#,
+        "</invoke></minimax:tool_call>",
+    ];
+
+    let mut second_tool_name = None;
+    for chunk in chunks2 {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        for call in result.calls {
+            if let Some(name) = call.name {
+                second_tool_name = Some(name);
+            }
+        }
+    }
+
+    assert_eq!(second_tool_name, Some("search".to_string()));
+}
+
+#[tokio::test]
+async fn test_minimax_many_parameters() {
+    let parser = MinimaxM2Parser::new();
+
+    let mut params_xml = String::new();
+    for i in 1..=20 {
+        params_xml.push_str(&format!(
+            r#"<parameter name="param{}">value{}</parameter>
+"#,
+            i, i
+        ));
+    }
+
+    let input = format!(
+        r#"<minimax:tool_call>
+<invoke name="complex_func">
+{}
+</invoke>
+</minimax:tool_call>"#,
+        params_xml
+    );
+
+    let (_normal_text, tools) = parser.parse_complete(&input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "complex_func");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+
+    // Verify all 20 parameters are parsed
+    for i in 1..=20 {
+        let key = format!("param{}", i);
+        let expected_value = format!("value{}", i);
+        assert_eq!(args[key], expected_value);
+    }
+}
+
+#[tokio::test]
+async fn test_minimax_character_by_character_streaming() {
+    // Test character-by-character streaming to simulate real-world streaming
+    let mut parser = MinimaxM2Parser::new();
+    let tools = create_test_tools();
+
+    let complete_text = r#"Let me help you. <minimax:tool_call>
+<invoke name="get_weather">
+<parameter name="city">Seattle</parameter>
+<parameter name="units">celsius</parameter>
+</invoke>
+</minimax:tool_call> Here are the results."#;
+
+    let mut content_collected = String::new();
+    let mut tool_name_found = false;
+    let mut parameters_found = Vec::new();
+
+    // Stream character by character - feed only one character at a time
+    for i in 0..complete_text.len() {
+        let delta = &complete_text[i..i + 1];
+        let result = parser.parse_incremental(delta, &tools).await.unwrap();
+        content_collected.push_str(&result.normal_text);
+
+        for call in result.calls {
+            if let Some(name) = call.name {
+                assert_eq!(name, "get_weather");
+                tool_name_found = true;
+            }
+            if !call.parameters.is_empty() && !parameters_found.contains(&call.parameters) {
+                parameters_found.push(call.parameters.clone());
+            }
+        }
+    }
+
+    assert!(
+        tool_name_found,
+        "Should find tool name during character-by-character streaming"
+    );
+    assert!(
+        !parameters_found.is_empty(),
+        "Should find parameters during streaming"
+    );
+
+    // Should have initial content and final content
+    assert!(content_collected.contains("Let me help you."));
+    assert!(content_collected.contains("Here are the results."));
+}
+
+#[tokio::test]
+async fn test_minimax_content_before_and_after_tool_calls() {
+    let parser = MinimaxM2Parser::new();
+
+    let input = r#"I'll analyze the weather for you now.
+<minimax:tool_call>
+<invoke name="get_weather">
+<parameter name="city">Boston</parameter>
+<parameter name="state">MA</parameter>
+</invoke>
+</minimax:tool_call>
+Based on the analysis, here's what I found."#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+
+    // Verify tool extraction
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "get_weather");
+
+    // Verify content preservation (only text before tool call is returned)
+    assert!(normal_text.contains("I'll analyze the weather for you now."));
+    // Text after tool call is not included in parse_complete
+    assert!(!normal_text.contains("Based on the analysis, here's what I found."));
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["city"], "Boston");
+    assert_eq!(args["state"], "MA");
+}
+
+#[tokio::test]
+async fn test_minimax_incomplete_tool_call() {
+    let parser = MinimaxM2Parser::new();
+
+    // Incomplete tool call - missing closing tag
+    let input = r#"<minimax:tool_call>
+<invoke name="get_weather">
+<parameter name="city">Chicago</parameter>"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+
+    // Should not extract incomplete tool calls
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input); // Should return as normal text
+}
+
+#[tokio::test]
+async fn test_minimax_malformed_invoke_tag() {
+    let parser = MinimaxM2Parser::new();
+
+    // Malformed invoke tag - missing name attribute
+    let input = r#"<minimax:tool_call>
+<invoke>
+<parameter name="city">Miami</parameter>
+</invoke>
+</minimax:tool_call>"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+
+    // Should not extract tool calls with malformed invoke tags
+    assert_eq!(tools.len(), 0);
+    assert_eq!(normal_text, input);
+}
+
+#[tokio::test]
+async fn test_minimax_streaming_with_invalid_function_progressive() {
+    let mut parser = MinimaxM2Parser::new();
+    let tools = create_test_tools();
+
+    // Progressive chunks building an invalid function call
+    let chunks = vec![
+        "<minimax:tool_call>",
+        r#"<invoke name="invalid_function">"#,
+        r#"<parameter name="test">value</parameter>"#,
+        "</invoke>",
+        "</minimax:tool_call>",
+    ];
+
+    let mut all_normal_text = String::new();
+    let mut found_valid_tool = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        all_normal_text.push_str(&result.normal_text);
+
+        for call in result.calls {
+            if let Some(name) = call.name {
+                // Should not get here for invalid function
+                if tools.iter().any(|t| t.function.name == name) {
+                    found_valid_tool = true;
+                }
+            }
+        }
+    }
+
+    assert!(
+        !found_valid_tool,
+        "Invalid function should not be parsed as tool call"
+    );
+    // The invalid tool call should be returned as normal text
+    assert!(all_normal_text.contains("invalid_function"));
+}
+
+#[tokio::test]
+async fn test_minimax_rapid_streaming_bursts() {
+    // Test handling of rapid streaming bursts (multiple chunks at once)
+    let mut parser = MinimaxM2Parser::new();
+    let tools = create_test_tools();
+
+    let chunks = vec![
+        "<minimax:tool_call><invoke name=\"search\"><parameter name=\"query\">",
+        "rust programming",
+        "</parameter></invoke></minimax:tool_call>",
+    ];
+
+    let mut found_function = false;
+    let mut parameters = Vec::new();
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+
+        for call in result.calls {
+            if let Some(name) = call.name {
+                assert_eq!(name, "search");
+                found_function = true;
+            }
+            if !call.parameters.is_empty() {
+                parameters.push(call.parameters.clone());
+            }
+        }
+    }
+
+    assert!(found_function);
+
+    // Verify that parameters were streamed correctly
+    let final_params = parameters.join("");
+    assert!(final_params.contains("rust programming"));
+}
+
+#[tokio::test]
+async fn test_minimax_special_characters_in_values() {
+    let parser = MinimaxM2Parser::new();
+
+    let input = r#"<minimax:tool_call>
+<invoke name="process">
+<parameter name="text">Special chars: @#$%^&*()</parameter>
+<parameter name="emoji">🦀 Rust 🚀</parameter>
+<parameter name="quotes">"double" and 'single' quotes</parameter>
+</invoke>
+</minimax:tool_call>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["text"], "Special chars: @#$%^&*()");
+    assert_eq!(args["emoji"], "🦀 Rust 🚀");
+    assert_eq!(args["quotes"], "\"double\" and 'single' quotes");
+}
+
+#[tokio::test]
+async fn test_minimax_whitespace_handling() {
+    let parser = MinimaxM2Parser::new();
+
+    // Test with various whitespace scenarios
+    let input = r#"<minimax:tool_call>
+    <invoke name="process">
+        <parameter name="trimmed">  spaces around  </parameter>
+        <parameter name="newlines">
+            Line 1
+            Line 2
+        </parameter>
+        <parameter name="tabs">	tab	separated	</parameter>
+    </invoke>
+</minimax:tool_call>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    // Values should preserve internal whitespace but may trim edges based on parser design
+    assert!(args["newlines"].as_str().unwrap().contains("Line 1"));
+    assert!(args["newlines"].as_str().unwrap().contains("Line 2"));
+    assert_eq!(args["tabs"], "\ttab\tseparated\t");
+}
+
+#[tokio::test]
+async fn test_minimax_no_tools() {
+    // Test input with no tool calls at all
+    let parser = MinimaxM2Parser::new();
+
+    let input = r#"This is just a normal response without any tool calls.
+I can provide information directly without using any tools.
+Even if I mention function names like get_weather or search,
+they are not actual tool calls unless properly formatted."#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+
+    // No tools should be extracted
+    assert_eq!(
+        tools.len(),
+        0,
+        "Should not extract any tools from plain text"
+    );
+
+    // All content should be returned as normal text
+    assert_eq!(
+        normal_text, input,
+        "All content should be returned as normal text when no tools present"
+    );
+}
+
+#[tokio::test]
+async fn test_minimax_invalid_json_in_parameters() {
+    // Test handling of invalid JSON in parameter values
+    let parser = MinimaxM2Parser::new();
+
+    let input = r#"<minimax:tool_call>
+<invoke name="process">
+<parameter name="valid">{"key": "value"}</parameter>
+<parameter name="invalid">{invalid json: no quotes}</parameter>
+<parameter name="broken">[1, 2, unclosed</parameter>
+<parameter name="mixed">Some text {"partial": json} more text</parameter>
+</invoke>
+</minimax:tool_call>"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+
+    // Tool should still be extracted despite invalid JSON in parameters
+    assert_eq!(
+        tools.len(),
+        1,
+        "Should extract tool even with invalid JSON in parameters"
+    );
+    assert_eq!(tools[0].function.name, "process");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+
+    // Parameters are stored as strings, not parsed as JSON
+    // Even invalid JSON should be preserved as string values
+    assert!(args["valid"].is_string());
+    assert_eq!(args["valid"], r#"{"key": "value"}"#);
+
+    assert!(args["invalid"].is_string());
+    assert_eq!(args["invalid"], "{invalid json: no quotes}");
+
+    assert!(args["broken"].is_string());
+    assert_eq!(args["broken"], "[1, 2, unclosed");
+
+    assert!(args["mixed"].is_string());
+    assert_eq!(args["mixed"], r#"Some text {"partial": json} more text"#);
+
+    assert_eq!(normal_text, "");
+}

sgl-model-gateway/tests/tool_parser_mistral.rs

@@ -0,0 +1,274 @@
+//! Mistral Parser Integration Tests
+//!
+//! Tests for the Mistral parser which handles [TOOL_CALLS] format
+
+use serde_json::json;
+use sgl_model_gateway::tool_parser::{MistralParser, ToolParser};
+
+#[tokio::test]
+async fn test_mistral_single_tool() {
+    let parser = MistralParser::new();
+    let input = r#"Let me search for that.
+[TOOL_CALLS] [{"name": "search_web", "arguments": {"query": "latest news", "max_results": 5}}]"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "Let me search for that.\n");
+    assert_eq!(tools[0].function.name, "search_web");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["query"], "latest news");
+    assert_eq!(args["max_results"], 5);
+}
+
+#[tokio::test]
+async fn test_mistral_multiple_tools() {
+    let parser = MistralParser::new();
+    let input = r#"I'll help you with both tasks.
+[TOOL_CALLS] [
+    {"name": "get_weather", "arguments": {"city": "Tokyo", "units": "celsius"}},
+    {"name": "search_news", "arguments": {"query": "AI developments", "limit": 10}}
+]"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 2);
+    assert_eq!(normal_text, "I'll help you with both tasks.\n");
+
+    assert_eq!(tools[0].function.name, "get_weather");
+    let args0: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args0["city"], "Tokyo");
+
+    assert_eq!(tools[1].function.name, "search_news");
+    let args1: serde_json::Value = serde_json::from_str(&tools[1].function.arguments).unwrap();
+    assert_eq!(args1["query"], "AI developments");
+}
+
+#[tokio::test]
+async fn test_mistral_nested_json() {
+    let parser = MistralParser::new();
+    let input = r#"Processing complex data.
+[TOOL_CALLS] [{"name": "process_data", "arguments": {"config": {"nested": {"value": [1, 2, 3]}}, "enabled": true}}]"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "Processing complex data.\n");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["config"]["nested"]["value"], json!([1, 2, 3]));
+    assert_eq!(args["enabled"], true);
+}
+
+#[tokio::test]
+async fn test_mistral_with_text_after() {
+    let parser = MistralParser::new();
+    let input = r#"[TOOL_CALLS] [{"name": "test", "arguments": {}}]
+
+And here's some text after the tool call that should be ignored."#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "test");
+}
+
+#[tokio::test]
+async fn test_mistral_empty_arguments() {
+    let parser = MistralParser::new();
+    let input = r#"[TOOL_CALLS] [{"name": "ping", "arguments": {}}]"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "ping");
+}
+
+#[tokio::test]
+async fn test_mistral_with_brackets_in_strings() {
+    let parser = MistralParser::new();
+    let input = r#"[TOOL_CALLS] [{"name": "echo", "arguments": {"text": "Array notation: arr[0] = value[1]"}}]"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["text"], "Array notation: arr[0] = value[1]");
+}
+
+#[tokio::test]
+async fn test_mistral_format_detection() {
+    let parser = MistralParser::new();
+
+    assert!(parser.has_tool_markers("[TOOL_CALLS] ["));
+    assert!(parser.has_tool_markers("Some text [TOOL_CALLS] ["));
+    assert!(!parser.has_tool_markers("Just plain text"));
+    assert!(!parser.has_tool_markers("[{\"name\": \"test\"}]")); // JSON array without TOOL_CALLS
+}
+
+#[tokio::test]
+async fn test_mistral_malformed_json() {
+    let parser = MistralParser::new();
+
+    // Missing closing bracket
+    let input = r#"[TOOL_CALLS] [{"name": "test", "arguments": {}"#;
+    if let Ok((_normal_text, tools)) = parser.parse_complete(input).await {
+        assert_eq!(tools.len(), 0);
+    }
+    // Error is also acceptable for malformed input
+
+    // Invalid JSON inside
+    let input = r#"[TOOL_CALLS] [{"name": invalid}]"#;
+    if let Ok((_normal_text, tools)) = parser.parse_complete(input).await {
+        assert_eq!(tools.len(), 0);
+    }
+    // Error is also acceptable for malformed input
+}
+
+#[tokio::test]
+async fn test_mistral_real_world_output() {
+    let parser = MistralParser::new();
+
+    // Actual output from Mistral model
+    let input = r#"I'll search for information about Rust programming and check the weather in San Francisco.
+
+[TOOL_CALLS] [
+    {
+        "name": "web_search",
+        "arguments": {
+            "query": "Rust programming language features 2024",
+            "max_results": 3,
+            "include_snippets": true
+        }
+    },
+    {
+        "name": "get_weather",
+        "arguments": {
+            "location": "San Francisco, CA",
+            "units": "fahrenheit",
+            "include_forecast": false
+        }
+    }
+]
+
+Let me execute these searches for you."#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 2);
+    assert_eq!(normal_text, "I'll search for information about Rust programming and check the weather in San Francisco.\n\n");
+    assert_eq!(tools[0].function.name, "web_search");
+    assert_eq!(tools[1].function.name, "get_weather");
+}
+
+#[tokio::test]
+async fn test_mistral_streaming_closing_bracket() {
+    use sgl_model_gateway::protocols::common::Tool;
+
+    // Test that closing ] is stripped for Mistral array format
+    let mut parser = MistralParser::new();
+
+    let tools = vec![Tool {
+        tool_type: "function".to_string(),
+        function: sgl_model_gateway::protocols::common::Function {
+            name: "get_weather".to_string(),
+            description: Some("Get weather".to_string()),
+            parameters: json!({}),
+            strict: None,
+        },
+    }];
+
+    let chunks = vec![
+        "[TOOL_CALLS] ",
+        "[{",
+        "\"",
+        "name",
+        "\":",
+        "\"",
+        "get",
+        "_weather",
+        "\",",
+        "\"",
+        "arguments",
+        "\":",
+        "{",
+        "\"",
+        "city",
+        "\":",
+        "\"",
+        "Paris",
+        "\"",
+        "}",
+        "}",
+        "]",
+        " Here's",
+        " the weather",
+        " info",
+    ];
+
+    let mut all_normal_text = String::new();
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        all_normal_text.push_str(&result.normal_text);
+    }
+
+    // Should emit only the third chunk as normal text, NOT the ]
+    assert_eq!(
+        all_normal_text, " Here's the weather info",
+        "Should not emit ] for Mistral array format, got: '{}'",
+        all_normal_text
+    );
+}
+
+#[tokio::test]
+async fn test_mistral_streaming_bracket_in_text_after_tools() {
+    use sgl_model_gateway::protocols::common::Tool;
+
+    // Test that ] in normal text AFTER tool calls is preserved
+    let mut parser = MistralParser::new();
+
+    let tools = vec![Tool {
+        tool_type: "function".to_string(),
+        function: sgl_model_gateway::protocols::common::Function {
+            name: "get_weather".to_string(),
+            description: Some("Get weather".to_string()),
+            parameters: json!({}),
+            strict: None,
+        },
+    }];
+
+    let chunks = vec![
+        "[TOOL_CALLS] ",
+        "[",
+        "{",
+        "\"name",
+        "\":",
+        "\"get_weather",
+        "\",",
+        "\"arguments",
+        "\":",
+        "{\"",
+        "city",
+        "\":",
+        "\"Paris",
+        "\"}",
+        "}",
+        "]",
+        " Array",
+        " notation:",
+        " arr",
+        "[",
+        "0",
+        "]",
+    ];
+
+    let mut all_normal_text = String::new();
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        all_normal_text.push_str(&result.normal_text);
+    }
+
+    // Should preserve ] in normal text after tools complete
+    assert_eq!(
+        all_normal_text, " Array notation: arr[0]",
+        "Should preserve ] in normal text after tools, got: '{}'",
+        all_normal_text
+    );
+}

sgl-model-gateway/tests/tool_parser_mixed_edge_cases.rs

@@ -0,0 +1,291 @@
+//! Mixed Format and Additional Edge Case Tests
+//!
+//! Tests for edge cases across parsers and mixed format scenarios
+
+use serde_json::json;
+use sgl_model_gateway::tool_parser::{
+    JsonParser, LlamaParser, MistralParser, PythonicParser, QwenParser, ToolParser,
+};
+
+mod common;
+use common::create_test_tools;
+
+#[tokio::test]
+async fn test_mixed_formats_in_text() {
+    let json_parser = JsonParser::new();
+    let input = r#"
+    Some text with [TOOL_CALLS] marker that shouldn't trigger.
+    Also has <tool_call> tags and [function()] syntax.
+    But here's the actual JSON: {"name": "test", "arguments": {}}
+    "#;
+
+    let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "test");
+
+    // Mistral parser should ignore JSON and other formats
+    let mistral_parser = MistralParser::new();
+    let input = r#"
+    {"name": "fake"} [function()] <tool_call>
+    [TOOL_CALLS] [{"name": "real", "arguments": {}}]
+    "#;
+
+    let (_normal_text, tools) = mistral_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "real");
+}
+
+#[tokio::test]
+async fn test_format_markers_in_string_content() {
+    let pythonic_parser = PythonicParser::new();
+    let input = r#"[echo(text="Use [TOOL_CALLS] and <tool_call> in text")]"#;
+
+    let (_normal_text, tools) = pythonic_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["text"], "Use [TOOL_CALLS] and <tool_call> in text");
+
+    let qwen_parser = QwenParser::new();
+    let input = r#"<tool_call>
+{"name": "log", "arguments": {"msg": "Found [function()] pattern"}}
+</tool_call>"#;
+
+    let (_normal_text, tools) = qwen_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["msg"], "Found [function()] pattern");
+}
+
+#[tokio::test]
+async fn test_deeply_nested_json_structures() {
+    let json_parser = JsonParser::new();
+
+    let input = r#"{
+        "name": "deep_process",
+        "arguments": {
+            "level1": {
+                "level2": {
+                    "level3": {
+                        "level4": {
+                            "level5": {
+                                "data": [1, 2, [3, [4, 5]]]
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }"#;
+
+    let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "deep_process");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert!(args["level1"]["level2"]["level3"]["level4"]["level5"]["data"].is_array());
+}
+
+#[tokio::test]
+async fn test_multiple_sequential_calls_different_formats() {
+    // Simulate a scenario where different parts of text have different formats
+    // (though each parser will only recognize its own format)
+
+    let llama_parser = LlamaParser::new();
+
+    // Llama parser currently only returns the first tool found
+    let input = r#"First call: <|python_tag|>{"name": "call1", "arguments": {}}"#;
+
+    let (_normal_text, tools) = llama_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "call1");
+
+    let input2 = r#"{"name": "call2", "arguments": {"x": 1}}"#;
+    let (_normal_text2, tools2) = llama_parser.parse_complete(input2).await.unwrap();
+    assert_eq!(tools2.len(), 1);
+    assert_eq!(tools2[0].function.name, "call2");
+}
+
+#[tokio::test]
+async fn test_empty_and_whitespace_variations() {
+    let json_parser = JsonParser::new();
+
+    // Various whitespace scenarios
+    let cases = vec![
+        r#"  {"name":"compact","arguments":{}}  "#,
+        r#"
+
+        {"name": "spaced", "arguments": {}}
+
+        "#,
+        r#"	{"name": "tabbed", "arguments": {}}	"#, // tabs
+    ];
+
+    for input in cases {
+        let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
+        assert_eq!(tools.len(), 1, "Should parse regardless of whitespace");
+    }
+}
+
+#[tokio::test]
+async fn test_special_json_values() {
+    let json_parser = JsonParser::new();
+
+    let input = r#"{
+        "name": "test_special",
+        "arguments": {
+            "float_e": 1.23e10,
+            "float_neg_e": 1.23e-10,
+            "hex_like": "0x1234",
+            "very_long_num": 99999999999999999999,
+            "special_strings": ["", " ", "\u0000", "\u001f"],
+            "escaped": "\\n\\r\\t\\\"\\\\",
+            "unicode": "\u4e2d\u6587"
+        }
+    }"#;
+
+    let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "test_special");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert!(args["special_strings"].is_array());
+    assert!(args["escaped"].is_string());
+}
+
+#[tokio::test]
+async fn test_parser_recovery_after_invalid_input() {
+    let mut parser = JsonParser::new();
+    let tools = create_test_tools();
+
+    // Send invalid JSON first
+    let _ = parser.parse_incremental(r#"{"broken": "#, &tools).await;
+
+    // Create a new parser instance for clean state
+    let mut parser2 = JsonParser::new();
+    let result = parser2
+        .parse_incremental(r#"{"name": "valid", "arguments": {}}"#, &tools)
+        .await
+        .unwrap();
+
+    if !result.calls.is_empty() {
+        if let Some(name) = &result.calls[0].name {
+            assert_eq!(name, "valid");
+        }
+    }
+}
+
+#[tokio::test]
+async fn test_boundary_cases_for_extraction() {
+    let json_parser = JsonParser::new();
+
+    // JSON at the very beginning
+    let input = r#"{"name": "start", "arguments": {}} and then text"#;
+    let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "start");
+
+    // JSON at the very end
+    let input = r#"Some text first {"name": "end", "arguments": {}}"#;
+    let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "end");
+
+    // Multiple JSON objects in text (should find first valid one)
+    let input =
+        r#"Text {"name": "first", "arguments": {}} more {"name": "second", "arguments": {}}"#;
+    let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
+    assert!(!tools.is_empty());
+    assert_eq!(tools[0].function.name, "first");
+}
+
+#[tokio::test]
+async fn test_pythonic_edge_cases() {
+    let parser = PythonicParser::new();
+
+    // Function name with underscores and numbers
+    let input = r#"[func_name_2(param_1="value")]"#;
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "func_name_2");
+
+    // Empty string argument
+    let input = r#"[process(text="")]"#;
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["text"], "");
+}
+
+#[tokio::test]
+async fn test_mistral_with_pretty_json() {
+    let parser = MistralParser::new();
+
+    // Pretty-printed JSON in Mistral format
+    let input = r#"[TOOL_CALLS] [
+        {
+            "name": "formatted",
+            "arguments": {
+                "nested": {
+                    "key": "value"
+                },
+                "array": [
+                    1,
+                    2,
+                    3
+                ]
+            }
+        }
+    ]"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "formatted");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["nested"]["key"], "value");
+    assert_eq!(args["array"], json!([1, 2, 3]));
+}
+
+#[tokio::test]
+async fn test_qwen_with_cdata_like_content() {
+    let parser = QwenParser::new();
+
+    // Note: QwenParser expects exactly "<tool_call>\n" with the newline
+    let input = r#"<tool_call>
+{"name": "process", "arguments": {"xml": "<![CDATA[some data]]>"}}
+</tool_call>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "process");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["xml"], "<![CDATA[some data]]>");
+}
+
+#[tokio::test]
+async fn test_extremely_long_function_names() {
+    let parser = PythonicParser::new();
+
+    let long_name = "very_long_function_name_that_might_appear_in_generated_code_somewhere";
+    let input = format!(r#"[{}(param="value")]"#, long_name);
+
+    let (_normal_text, tools) = parser.parse_complete(&input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, long_name);
+}
+
+#[tokio::test]
+async fn test_json_with_duplicate_keys() {
+    let parser = JsonParser::new();
+
+    // JSON with duplicate keys (last one should win per JSON spec)
+    let input = r#"{"name": "test", "arguments": {"key": "first", "key": "second"}}"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    // JSON parsers typically keep the last value for duplicate keys
+    assert_eq!(args["key"], "second");
+}

sgl-model-gateway/tests/tool_parser_partial_json.rs

@@ -0,0 +1,156 @@
+//! Partial JSON Parser Tests
+//!
+//! Tests for the partial JSON parser with allow_partial_strings flag behavior
+
+use sgl_model_gateway::tool_parser::partial_json::PartialJson;
+
+#[test]
+fn test_partial_string_flag_disallows_incomplete_strings() {
+    // Test case from the bug report: {"name": "
+    // With allow_partial_strings=false, should return {} (stop before incomplete string)
+    let parser = PartialJson::new(32, true);
+    let input = r#"{"name": ""#;
+
+    let result = parser.parse_value(input, false);
+    assert!(result.is_ok());
+
+    let (obj, consumed) = result.unwrap();
+
+    // Should parse just the opening brace and stop at the incomplete string
+    assert!(obj.is_object());
+    let obj_map = obj.as_object().unwrap();
+
+    // Should have empty object (stopped before parsing incomplete "name" key)
+    assert!(
+        obj_map.is_empty() || !obj_map.contains_key("name"),
+        "Should not parse incomplete string key, got: {:?}",
+        obj_map
+    );
+
+    // Should consume characters up to the incomplete string
+    assert!(consumed <= input.len());
+}
+
+#[test]
+fn test_partial_string_flag_allows_incomplete_strings() {
+    // Test case: {"name": "
+    // With allow_partial_strings=true, should parse the incomplete string
+    let parser = PartialJson::new(32, true);
+    let input = r#"{"name": ""#;
+
+    let result = parser.parse_value(input, true);
+    assert!(result.is_ok());
+
+    let (obj, consumed) = result.unwrap();
+
+    // Should parse the object with incomplete string value
+    assert!(obj.is_object());
+    let obj_map = obj.as_object().unwrap();
+
+    // With allow_partial_strings=true, should parse "name" key with empty string value
+    assert!(
+        obj_map.contains_key("name"),
+        "Should parse incomplete string with allow_partial_strings=true"
+    );
+
+    assert_eq!(consumed, input.len());
+}
+
+#[test]
+fn test_partial_string_flag_complete_json() {
+    // Test case: {"name": "test"}
+    // Both flags should parse complete JSON the same way
+    let input = r#"{"name": "test"}"#;
+
+    let parser = PartialJson::new(32, true);
+    let result1 = parser.parse_value(input, false);
+    assert!(result1.is_ok());
+    let (obj1, consumed1) = result1.unwrap();
+
+    let result2 = parser.parse_value(input, true);
+    assert!(result2.is_ok());
+    let (obj2, consumed2) = result2.unwrap();
+
+    // Both should parse the same complete JSON
+    assert_eq!(obj1, obj2);
+    assert_eq!(consumed1, consumed2);
+    assert_eq!(consumed1, input.len());
+
+    // Check the parsed value
+    assert!(obj1.is_object());
+    let obj_map = obj1.as_object().unwrap();
+    assert_eq!(obj_map.get("name").and_then(|v| v.as_str()), Some("test"));
+}
+
+#[test]
+fn test_backward_compatibility_default() {
+    // Test that default PartialJson still allows partial strings (backward compatible)
+    let parser = PartialJson::default();
+    let input = r#"{"name": ""#;
+
+    let result = parser.parse_value(input, true);
+    assert!(result.is_ok());
+
+    let (obj, _) = result.unwrap();
+    assert!(obj.is_object());
+
+    // Default behavior should allow partial strings
+    let obj_map = obj.as_object().unwrap();
+    assert!(
+        obj_map.contains_key("name"),
+        "Default should allow partial strings for backward compatibility"
+    );
+}
+
+#[test]
+fn test_partial_string_in_nested_object() {
+    // Test case: {"tool": {"name": "
+    let parser = PartialJson::new(32, true);
+    let input = r#"{"tool": {"name": ""#;
+
+    let result = parser.parse_value(input, false);
+    assert!(result.is_ok());
+
+    let (obj, _) = result.unwrap();
+    assert!(obj.is_object());
+
+    // With allow_partial_strings=false, should stop before incomplete nested string
+    let obj_map = obj.as_object().unwrap();
+    if let Some(tool) = obj_map.get("tool") {
+        if let Some(tool_map) = tool.as_object() {
+            assert!(
+                !tool_map.contains_key("name")
+                    || tool_map.get("name").and_then(|v| v.as_str()).is_none(),
+                "Should not parse incomplete nested string"
+            );
+        }
+    }
+}
+
+#[test]
+fn test_bug_fix_exact_scenario() {
+    // This test verifies the exact bug scenario from the issue:
+    // buffer = "{\"name\": \""
+    // flags = Allow.ALL & ~Allow.STR
+    // Python returns: Parsed object: {}, consumed length: 10
+
+    let parser = PartialJson::new(32, true);
+    let input = r#"{"name": ""#;
+
+    let result = parser.parse_value(input, false);
+    assert!(result.is_ok());
+
+    let (obj, consumed) = result.unwrap();
+
+    // Should return empty object (not {"name": null} or {"name": ""})
+    assert!(obj.is_object());
+    let obj_map = obj.as_object().unwrap();
+    assert!(
+        obj_map.is_empty(),
+        "Expected empty object, got: {:?}. This matches Python behavior with Allow.ALL & ~Allow.STR",
+        obj_map
+    );
+
+    // Should consume all characters (10 bytes)
+    assert_eq!(consumed, 10, "Should consume all 10 characters");
+}

sgl-model-gateway/tests/tool_parser_pythonic.rs

@@ -0,0 +1,518 @@
+//! Pythonic Parser Integration Tests
+//!
+//! Tests for the Pythonic parser which handles Python function call syntax
+
+use serde_json::json;
+use sgl_model_gateway::tool_parser::{PythonicParser, ToolParser};
+
+mod common;
+use common::create_test_tools;
+
+#[tokio::test]
+async fn test_pythonic_single_function() {
+    let parser = PythonicParser::new();
+    let input = r#"[get_weather(city="London", units="celsius")]"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "get_weather");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["city"], "London");
+    assert_eq!(args["units"], "celsius");
+}
+
+#[tokio::test]
+async fn test_pythonic_multiple_functions() {
+    let parser = PythonicParser::new();
+    let input =
+        r#"[search_web(query="Rust programming", max_results=5), get_time(timezone="UTC")]"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 2);
+    assert_eq!(tools[0].function.name, "search_web");
+    assert_eq!(tools[1].function.name, "get_time");
+
+    let args0: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args0["query"], "Rust programming");
+    assert_eq!(args0["max_results"], 5);
+}
+
+#[tokio::test]
+async fn test_pythonic_with_python_literals() {
+    let parser = PythonicParser::new();
+    let input = r#"[configure(enabled=True, disabled=False, optional=None)]"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["enabled"], true);
+    assert_eq!(args["disabled"], false);
+    assert_eq!(args["optional"], json!(null));
+}
+
+#[tokio::test]
+async fn test_pythonic_with_lists_and_dicts() {
+    let parser = PythonicParser::new();
+    let input =
+        r#"[process_data(items=[1, 2, 3], config={"key": "value", "nested": {"deep": True}})]"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["items"], json!([1, 2, 3]));
+    assert_eq!(args["config"]["key"], "value");
+    assert_eq!(args["config"]["nested"]["deep"], true);
+}
+
+#[tokio::test]
+async fn test_pythonic_with_special_tokens() {
+    let parser = PythonicParser::new();
+
+    // Llama 4 sometimes outputs these tokens
+    let input = r#"<|python_start|>[calculate(x=10, y=20)]<|python_end|>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "calculate");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["x"], 10);
+    assert_eq!(args["y"], 20);
+}
+
+#[tokio::test]
+async fn test_pythonic_with_nested_parentheses() {
+    let parser = PythonicParser::new();
+    let input = r#"[math_eval(expression="(2 + 3) * (4 - 1)", round_to=2)]"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["expression"], "(2 + 3) * (4 - 1)");
+    assert_eq!(args["round_to"], 2);
+}
+
+#[tokio::test]
+async fn test_pythonic_with_escaped_quotes() {
+    let parser = PythonicParser::new();
+    let input = r#"[echo(text="She said \"Hello\" to him")]"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["text"], "She said \"Hello\" to him");
+}
+
+#[tokio::test]
+async fn test_pythonic_empty_arguments() {
+    let parser = PythonicParser::new();
+    let input = r#"[ping()]"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "ping");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args, json!({}));
+}
+
+#[tokio::test]
+async fn test_pythonic_format_detection() {
+    let parser = PythonicParser::new();
+
+    assert!(!parser.has_tool_markers("[function_name(")); // Incomplete
+    assert!(parser.has_tool_markers("[get_weather(city=\"NYC\")]"));
+    assert!(!parser.has_tool_markers("Just plain text"));
+    assert!(!parser.has_tool_markers("{\"name\": \"test\"}")); // JSON
+}
+
+#[tokio::test]
+async fn test_pythonic_invalid_syntax() {
+    let parser = PythonicParser::new();
+
+    // Missing closing bracket
+    let input = r#"[function(arg=value"#;
+    if let Ok((_normal_text, tools)) = parser.parse_complete(input).await {
+        assert_eq!(tools.len(), 0);
+    }
+    // Error is also acceptable for invalid syntax
+
+    // Invalid Python syntax - empty parameter name
+    // Note: The parser currently accepts this invalid syntax and returns a result
+    // This is a known limitation of the current implementation
+    let input = r#"[function(=value)]"#;
+    if let Ok((_normal_text, tools)) = parser.parse_complete(input).await {
+        // The parser incorrectly accepts this, returning 1 result
+        // We'll accept this behavior for now but note it's not ideal
+        assert!(tools.len() <= 1, "Should parse at most one function");
+    }
+    // Error would be the correct behavior
+}
+
+#[tokio::test]
+async fn test_pythonic_real_world_llama4() {
+    let parser = PythonicParser::new();
+
+    // Actual output from Llama 4 model
+    let input = r#"I'll help you with multiple tasks. Let me search for information and perform calculations.
+
+[web_search(query="latest Rust features", max_results=3, safe_search=True),
+ calculate(expression="42 * 3.14159", precision=2),
+ get_weather(city="San Francisco", units="fahrenheit", include_forecast=False)]
+
+These functions will provide the information you need."#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 3);
+    assert_eq!(normal_text, "I'll help you with multiple tasks. Let me search for information and perform calculations.\n\n\n\nThese functions will provide the information you need.");
+    assert_eq!(tools[0].function.name, "web_search");
+    assert_eq!(tools[1].function.name, "calculate");
+    assert_eq!(tools[2].function.name, "get_weather");
+
+    let args0: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args0["query"], "latest Rust features");
+    assert_eq!(args0["safe_search"], true);
+}
+
+#[tokio::test]
+async fn test_pythonic_nested_brackets_in_lists() {
+    let parser = PythonicParser::new();
+
+    let input = r#"[process_matrix(data=[[1, 2], [3, 4]], labels=["row[0]", "row[1]"])]"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "process_matrix");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["data"], json!([[1, 2], [3, 4]]));
+    assert_eq!(args["labels"], json!(["row[0]", "row[1]"]));
+}
+
+#[tokio::test]
+async fn test_pythonic_nested_brackets_in_dicts() {
+    let parser = PythonicParser::new();
+
+    let input =
+        r#"[analyze(config={"patterns": ["[a-z]+", "[0-9]+"], "nested": {"list": [1, [2, 3]]}})]"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "analyze");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["config"]["patterns"], json!(["[a-z]+", "[0-9]+"]));
+    assert_eq!(args["config"]["nested"]["list"], json!([1, [2, 3]]));
+}
+
+#[tokio::test]
+async fn test_pythonic_mixed_quotes() {
+    let parser = PythonicParser::new();
+
+    let input = r#"[format_text(single='Hello', double="World", mixed="It's \"quoted\"")]"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "format_text");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["single"], "Hello");
+    assert_eq!(args["double"], "World");
+    assert_eq!(args["mixed"], "It's \"quoted\"");
+}
+
+#[tokio::test]
+async fn test_pythonic_complex_nesting() {
+    let parser = PythonicParser::new();
+
+    let input = r#"[transform(
+        matrix=[[1, [2, 3]], [4, [5, [6, 7]]]],
+        operations=[{"type": "scale", "factor": [2, 3]}, {"type": "rotate", "angle": 90}],
+        metadata={"tags": ["nested[0]", "nested[1]"], "config": {"depth": [1, 2, 3]}}
+    )]"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "transform");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert!(args["matrix"].is_array());
+    assert!(args["operations"].is_array());
+    assert_eq!(args["operations"][0]["type"], "scale");
+    assert_eq!(args["metadata"]["config"]["depth"], json!([1, 2, 3]));
+}
+
+#[tokio::test]
+async fn test_parse_streaming_no_brackets() {
+    let mut parser = PythonicParser::new();
+
+    let tools = create_test_tools();
+
+    let text = "This is just normal text without any tool calls.";
+    let result = parser.parse_incremental(text, &tools).await.unwrap();
+
+    // Expected - no tool calls found
+    assert!(result.calls.is_empty());
+}
+
+#[tokio::test]
+async fn test_parse_streaming_complete_tool_call() {
+    let mut parser = PythonicParser::new();
+
+    let tools = create_test_tools();
+
+    let text = "Here's a tool call: [get_weather(location='New York', unit='celsius')]";
+    let result = parser.parse_incremental(text, &tools).await.unwrap();
+
+    assert!(!result.calls.is_empty(), "Should parse complete tool call");
+    assert_eq!(result.calls[0].name.as_ref().unwrap(), "get_weather");
+    let args: serde_json::Value = serde_json::from_str(&result.calls[0].parameters).unwrap();
+    assert_eq!(args["location"], "New York");
+    assert_eq!(args["unit"], "celsius");
+}
+
+#[tokio::test]
+async fn test_parse_streaming_text_before_tool_call() {
+    let mut parser = PythonicParser::new();
+
+    let tools = create_test_tools();
+
+    let text = "This is some text before [get_weather(location='London')]";
+    let result = parser.parse_incremental(text, &tools).await.unwrap();
+
+    assert!(!result.calls.is_empty(), "Should parse tool call");
+    assert_eq!(result.calls[0].name.as_ref().unwrap(), "get_weather");
+    let args: serde_json::Value = serde_json::from_str(&result.calls[0].parameters).unwrap();
+    assert_eq!(args["location"], "London");
+}
+
+#[tokio::test]
+async fn test_parse_streaming_partial_tool_call() {
+    let mut parser = PythonicParser::new();
+
+    let tools = create_test_tools();
+
+    // First chunk with opening bracket but no closing bracket
+    let text1 = "Let me check the weather: [get_weather(location=";
+    let result1 = parser.parse_incremental(text1, &tools).await.unwrap();
+
+    // First chunk should be incomplete
+    assert!(
+        result1.calls.is_empty(),
+        "First chunk should not return tool call"
+    );
+
+    // Second chunk completing the tool call
+    let text2 = "'Paris')]";
+    let result2 = parser.parse_incremental(text2, &tools).await.unwrap();
+
+    assert!(
+        !result2.calls.is_empty(),
+        "Second chunk should complete tool call"
+    );
+    assert_eq!(result2.calls[0].name.as_ref().unwrap(), "get_weather");
+    let args: serde_json::Value = serde_json::from_str(&result2.calls[0].parameters).unwrap();
+    assert_eq!(args["location"], "Paris");
+}
+
+#[tokio::test]
+async fn test_parse_streaming_bracket_without_text_before() {
+    let mut parser = PythonicParser::new();
+
+    let tools = create_test_tools();
+
+    let text = "[search(query='python programming')]";
+    let result = parser.parse_incremental(text, &tools).await.unwrap();
+
+    assert!(!result.calls.is_empty(), "Should parse tool call");
+    assert_eq!(result.calls[0].name.as_ref().unwrap(), "search");
+    let args: serde_json::Value = serde_json::from_str(&result.calls[0].parameters).unwrap();
+    assert_eq!(args["query"], "python programming");
+}
+
+#[tokio::test]
+async fn test_parse_streaming_text_after_tool_call() {
+    let mut parser = PythonicParser::new();
+
+    let tools = create_test_tools();
+
+    // First chunk with complete tool call and some text after
+    let text = "[get_weather(location='Tokyo')] Here's the forecast:";
+    let result = parser.parse_incremental(text, &tools).await.unwrap();
+
+    assert!(!result.calls.is_empty(), "Should parse tool call");
+    assert_eq!(result.calls[0].name.as_ref().unwrap(), "get_weather");
+    // Text after tool call is handled by parser internally
+}
+
+#[tokio::test]
+async fn test_parse_streaming_multiple_tool_calls() {
+    let mut parser = PythonicParser::new();
+
+    let tools = create_test_tools();
+
+    let text = "[get_weather(location='Berlin'), search(query='restaurants')]";
+
+    // Current implementation may handle this as a single parse
+    let result = parser.parse_incremental(text, &tools).await.unwrap();
+
+    // The parser should handle multiple tools in one bracket pair
+    // This test is flexible about the implementation behavior
+    if !result.calls.is_empty() {
+        // Parser found at least one tool
+        assert!(result.calls[0].name.is_some());
+    }
+    // Also acceptable if parser returns empty waiting for more context
+}
+
+#[tokio::test]
+async fn test_parse_streaming_opening_bracket_only() {
+    let mut parser = PythonicParser::new();
+
+    let tools = create_test_tools();
+
+    let text = "Let's try this: [";
+    let result = parser.parse_incremental(text, &tools).await.unwrap();
+
+    // Should be incomplete - no complete tool call
+    assert!(
+        result.calls.is_empty(),
+        "Should not return tool call for partial bracket"
+    );
+}
+
+#[tokio::test]
+async fn test_parse_streaming_nested_brackets() {
+    let mut parser = PythonicParser::new();
+
+    let tools = create_test_tools();
+
+    let text = "[get_weather(location='New York', unit='celsius', data=[1, 2, 3])]";
+    let result = parser.parse_incremental(text, &tools).await.unwrap();
+
+    assert!(
+        !result.calls.is_empty(),
+        "Should parse tool call with nested brackets"
+    );
+    assert_eq!(result.calls[0].name.as_ref().unwrap(), "get_weather");
+    let args: serde_json::Value = serde_json::from_str(&result.calls[0].parameters).unwrap();
+    assert_eq!(args["location"], "New York");
+    assert_eq!(args["unit"], "celsius");
+    assert_eq!(args["data"], json!([1, 2, 3]));
+}
+
+#[tokio::test]
+async fn test_parse_streaming_nested_brackets_dict() {
+    let mut parser = PythonicParser::new();
+    let tools = create_test_tools();
+
+    let text = r#"[search(query='test', config={'options': [1, 2], 'nested': {'key': 'value'}})]"#;
+    let result = parser.parse_incremental(text, &tools).await.unwrap();
+
+    assert!(
+        !result.calls.is_empty(),
+        "Should parse tool call with nested dict"
+    );
+    assert_eq!(result.calls[0].name.as_ref().unwrap(), "search");
+    let args: serde_json::Value = serde_json::from_str(&result.calls[0].parameters).unwrap();
+    assert_eq!(args["query"], "test");
+    assert_eq!(args["config"]["options"], json!([1, 2]));
+    assert_eq!(args["config"]["nested"]["key"], "value");
+}
+
+#[tokio::test]
+async fn test_parse_streaming_multiple_tools_with_nested_brackets() {
+    let mut parser = PythonicParser::new();
+
+    let tools = create_test_tools();
+
+    let text =
+        "[get_weather(location='Paris', data=[10, 20]), search(query='test', filters=['a', 'b'])]";
+    let result = parser.parse_incremental(text, &tools).await.unwrap();
+
+    // Should parse tools successfully
+    if !result.calls.is_empty() {
+        // At least gets the first tool
+        assert!(result.calls[0].name.is_some());
+    }
+}
+
+#[tokio::test]
+async fn test_parse_streaming_partial_nested_brackets() {
+    let mut parser = PythonicParser::new();
+
+    let tools = create_test_tools();
+
+    // First chunk with nested brackets but incomplete
+    let text1 = "Here's a call: [get_weather(location='Tokyo', data=[1, 2";
+    let result1 = parser.parse_incremental(text1, &tools).await.unwrap();
+
+    // First chunk should be incomplete
+    assert!(result1.calls.is_empty(), "First chunk should not complete");
+
+    // Second chunk completing the nested brackets
+    let text2 = ", 3])]";
+    let result2 = parser.parse_incremental(text2, &tools).await.unwrap();
+
+    assert!(
+        !result2.calls.is_empty(),
+        "Second chunk should complete tool call"
+    );
+    assert_eq!(result2.calls[0].name.as_ref().unwrap(), "get_weather");
+    let args: serde_json::Value = serde_json::from_str(&result2.calls[0].parameters).unwrap();
+    assert_eq!(args["location"], "Tokyo");
+    assert_eq!(args["data"], json!([1, 2, 3]));
+}
+
+#[tokio::test]
+async fn test_parse_streaming_with_python_start_and_end_token() {
+    let mut parser = PythonicParser::new();
+
+    let tools = create_test_tools();
+
+    let chunks = vec![
+        "Here's a call: ",
+        "<|python_",
+        "start|>[get_weather(location=",
+        "'Tokyo', data=[1, 2",
+        ", 3])]<|python_end|>",
+    ];
+
+    let mut got_tool = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        if !result.calls.is_empty() {
+            if let Some(name) = &result.calls[0].name {
+                assert_eq!(name, "get_weather");
+                let args: serde_json::Value =
+                    serde_json::from_str(&result.calls[0].parameters).unwrap();
+                assert_eq!(args["location"], "Tokyo");
+                assert_eq!(args["data"], json!([1, 2, 3]));
+                got_tool = true;
+            }
+        }
+    }
+
+    assert!(got_tool, "Should have parsed the tool call");
+}
+
+#[tokio::test]
+async fn test_detect_and_parse_with_python_start_and_end_token() {
+    let parser = PythonicParser::new();
+
+    let text = "User wants to get the weather in Mars. <|python_start|>[get_weather(location='Mars', unit='celsius')]<|python_end|> In this way we will get the weather in Mars.";
+    let (_normal_text, tools) = parser.parse_complete(text).await.unwrap();
+
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "get_weather");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["location"], "Mars");
+    assert_eq!(args["unit"], "celsius");
+}

sgl-model-gateway/tests/tool_parser_qwen.rs

@@ -0,0 +1,307 @@
+//! Qwen Parser Integration Tests
+//!
+//! Tests for the Qwen parser which handles <tool_call>...</tool_call> format
+
+use serde_json::json;
+use sgl_model_gateway::tool_parser::{QwenParser, ToolParser};
+
+mod common;
+use common::{create_test_tools, streaming_helpers::*};
+
+#[tokio::test]
+async fn test_qwen_single_tool() {
+    let parser = QwenParser::new();
+    let input = r#"<tool_call>
+{"name": "get_weather", "arguments": {"city": "Beijing", "units": "celsius"}}
+</tool_call>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "get_weather");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["city"], "Beijing");
+    assert_eq!(args["units"], "celsius");
+}
+
+#[tokio::test]
+async fn test_qwen_multiple_sequential_tools() {
+    let parser = QwenParser::new();
+    let input = r#"Let me help you with that.
+<tool_call>
+{"name": "search", "arguments": {"query": "Qwen model"}}
+</tool_call>
+<tool_call>
+{"name": "translate", "arguments": {"text": "Hello", "to": "zh"}}
+</tool_call>"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 2);
+    assert_eq!(normal_text, "Let me help you with that.\n");
+    assert_eq!(tools[0].function.name, "search");
+    assert_eq!(tools[1].function.name, "translate");
+}
+
+#[tokio::test]
+async fn test_qwen_pretty_printed_json() {
+    let parser = QwenParser::new();
+    let input = r#"<tool_call>
+{
+    "name": "create_document",
+    "arguments": {
+        "title": "Test Document",
+        "content": "This is a test",
+        "metadata": {
+            "author": "Qwen",
+            "tags": ["test", "example"]
+        }
+    }
+}
+</tool_call>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "create_document");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["metadata"]["author"], "Qwen");
+    assert_eq!(args["metadata"]["tags"], json!(["test", "example"]));
+}
+
+#[tokio::test]
+async fn test_qwen_with_text_between() {
+    let parser = QwenParser::new();
+    let input = r#"First, let me search for information.
+<tool_call>
+{"name": "search", "arguments": {"query": "test"}}
+</tool_call>
+
+Now I'll translate something.
+
+<tool_call>
+{"name": "translate", "arguments": {"text": "world", "to": "es"}}
+</tool_call>
+Done!"#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 2);
+    assert_eq!(normal_text, "First, let me search for information.\n");
+    assert_eq!(tools[0].function.name, "search");
+    assert_eq!(tools[1].function.name, "translate");
+}
+
+#[tokio::test]
+async fn test_qwen_empty_arguments() {
+    let parser = QwenParser::new();
+    let input = r#"<tool_call>
+{"name": "get_time", "arguments": {}}
+</tool_call>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "get_time");
+}
+
+#[tokio::test]
+async fn test_qwen_with_newlines_in_strings() {
+    let parser = QwenParser::new();
+    let input = r#"<tool_call>
+{"name": "write_file", "arguments": {"content": "Line 1\nLine 2\nLine 3", "path": "/tmp/test.txt"}}
+</tool_call>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["content"], "Line 1\nLine 2\nLine 3");
+}
+
+#[tokio::test]
+async fn test_qwen_format_detection() {
+    let parser = QwenParser::new();
+
+    assert!(parser.has_tool_markers("<tool_call>"));
+    assert!(parser.has_tool_markers("Some text <tool_call>\n{"));
+    assert!(!parser.has_tool_markers("Just plain text"));
+    assert!(!parser.has_tool_markers("{\"name\": \"test\"}")); // Plain JSON
+}
+
+#[tokio::test]
+async fn test_qwen_incomplete_tags() {
+    let parser = QwenParser::new();
+
+    // Missing closing tag
+    let input = r#"<tool_call>
+{"name": "test", "arguments": {}}"#;
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+
+    // Missing opening tag
+    let input = r#"{"name": "test", "arguments": {}}
+</tool_call>"#;
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0);
+}
+
+#[tokio::test]
+async fn test_qwen_real_world_output() {
+    let parser = QwenParser::new();
+
+    // Actual output from Qwen model
+    let input = r#"I'll help you search for information and perform calculations.
+
+<tool_call>
+{
+    "name": "web_search",
+    "arguments": {
+        "query": "quantum computing breakthroughs 2024",
+        "language": "en",
+        "region": "us",
+        "safe_search": true
+    }
+}
+</tool_call>
+
+Let me also calculate something for you:
+
+<tool_call>
+{
+    "name": "calculator",
+    "arguments": {
+        "expression": "sqrt(144) + 3^2",
+        "precision": 2
+    }
+}
+</tool_call>
+
+These tools will provide the information you need."#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 2);
+    assert_eq!(
+        normal_text,
+        "I'll help you search for information and perform calculations.\n\n"
+    );
+    assert_eq!(tools[0].function.name, "web_search");
+    assert_eq!(tools[1].function.name, "calculator");
+
+    let args0: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args0["query"], "quantum computing breakthroughs 2024");
+    assert_eq!(args0["safe_search"], true);
+}
+
+#[tokio::test]
+async fn test_buffer_drain_optimization() {
+    let mut parser = QwenParser::new();
+
+    let tools = create_test_tools();
+
+    // First chunk - incomplete tool call
+    let chunk1 = "<tool_call>\n{\"name\": \"test1\", ";
+    let _result = parser.parse_incremental(chunk1, &tools).await.unwrap();
+    // The important thing is buffer accumulation works
+
+    // Complete first tool and start second
+    let chunk2 = "\"arguments\": {}}\n</tool_call><tool_call>\n{\"name\": \"test2\", ";
+    let result = parser.parse_incremental(chunk2, &tools).await.unwrap();
+
+    if !result.calls.is_empty() {
+        if let Some(_name) = &result.calls[0].name {
+            assert_eq!(result.calls[0].name.as_ref().unwrap(), "test1");
+            // After consuming the first tool, buffer is managed internally
+        }
+    }
+
+    // Complete the second tool
+    let chunk3 = "\"arguments\": {\"x\": 1}}\n</tool_call>";
+    let result = parser.parse_incremental(chunk3, &tools).await.unwrap();
+
+    if !result.calls.is_empty() {
+        if let Some(_name) = &result.calls[0].name {
+            assert_eq!(result.calls[0].name.as_ref().unwrap(), "test2");
+            // Buffer is managed internally
+        }
+    }
+}
+
+#[tokio::test]
+async fn test_buffer_efficiency_with_multiple_tools() {
+    let mut parser = QwenParser::new();
+
+    let tools = create_test_tools();
+
+    // Send multiple complete tools at once
+    let input = r#"<tool_call>
+{"name": "tool1", "arguments": {"a": 1}}
+</tool_call><tool_call>
+{"name": "tool2", "arguments": {"b": 2}}
+</tool_call><tool_call>
+{"name": "tool3", "arguments": {"c": 3}}
+</tool_call>"#;
+
+    // This should efficiently process tools using drain() without creating new strings
+    let result = parser.parse_incremental(input, &tools).await.unwrap();
+
+    // In Phase 2, this will likely parse only the first tool
+    // The important thing is that drain() doesn't cause any issues
+    if !result.calls.is_empty() {
+        if let Some(name) = &result.calls[0].name {
+            assert!(["tool1", "tool2", "tool3"].contains(&name.as_str()));
+        }
+    }
+}
+
+// =============================================================================
+// REALISTIC STREAMING TESTS
+// =============================================================================
+
+#[tokio::test]
+async fn test_qwen_realistic_chunks_with_xml_tags() {
+    let tools = create_test_tools();
+    let mut parser = QwenParser::new();
+
+    let input = "<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Tokyo\"}}\n</tool_call>";
+    let chunks = create_realistic_chunks(input);
+
+    assert!(chunks.len() > 20, "Should have many small chunks");
+
+    let mut got_tool_name = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
+        for call in result.calls {
+            if let Some(name) = call.name {
+                assert_eq!(name, "get_weather");
+                got_tool_name = true;
+            }
+        }
+    }
+
+    assert!(got_tool_name, "Should have parsed tool name");
+}
+
+#[tokio::test]
+async fn test_qwen_xml_tag_arrives_in_parts() {
+    let tools = create_test_tools();
+    let mut parser = QwenParser::new();
+
+    let chunks = vec![
+        "<to", "ol_", "cal", "l>\n", "{", r#"""#, "na", "me", r#"""#, ": ", r#"""#, "tra", "nsl",
+        "ate", r#"""#, ", ", r#"""#, "arg", "ume", "nts", r#"""#, ": {", r#"""#, "tex", "t",
+        r#"""#, ": ", r#"""#, "hel", "lo", r#"""#, "}}\n", "</t", "ool", "_ca", "ll>",
+    ];
+
+    let mut got_tool_name = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+        for call in result.calls {
+            if let Some(name) = call.name {
+                assert_eq!(name, "translate");
+                got_tool_name = true;
+            }
+        }
+    }
+
+    assert!(got_tool_name, "Should have parsed tool name");
+}

sgl-model-gateway/tests/tool_parser_step3.rs

@@ -0,0 +1,238 @@
+//! Step3 Parser Integration Tests
+
+use sgl_model_gateway::tool_parser::{Step3Parser, ToolParser};
+
+mod common;
+use common::create_test_tools;
+
+#[tokio::test]
+async fn test_step3_complete_parsing() {
+    let parser = Step3Parser::new();
+
+    let input = r#"Let me help you.
+<｜tool_calls_begin｜>
+<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="search">
+<steptml:parameter name="query">rust programming</steptml:parameter>
+<steptml:parameter name="limit">10</steptml:parameter>
+</steptml:invoke><｜tool_call_end｜>
+<｜tool_calls_end｜>
+Here are the results..."#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "Let me help you.\n");
+    assert_eq!(tools[0].function.name, "search");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["query"], "rust programming");
+    assert_eq!(args["limit"], 10);
+}
+
+#[tokio::test]
+async fn test_step3_multiple_tools() {
+    let parser = Step3Parser::new();
+
+    let input = r#"<｜tool_calls_begin｜>
+<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="get_weather">
+<steptml:parameter name="location">Tokyo</steptml:parameter>
+</steptml:invoke><｜tool_call_end｜>
+<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="get_news">
+<steptml:parameter name="category">tech</steptml:parameter>
+<steptml:parameter name="limit">5</steptml:parameter>
+</steptml:invoke><｜tool_call_end｜>
+<｜tool_calls_end｜>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 2);
+    assert_eq!(tools[0].function.name, "get_weather");
+    assert_eq!(tools[1].function.name, "get_news");
+}
+
+#[tokio::test]
+async fn test_step3_type_conversion() {
+    let parser = Step3Parser::new();
+
+    let input = r#"<｜tool_calls_begin｜>
+<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="process">
+<steptml:parameter name="count">100</steptml:parameter>
+<steptml:parameter name="rate">2.5</steptml:parameter>
+<steptml:parameter name="active">true</steptml:parameter>
+<steptml:parameter name="optional">null</steptml:parameter>
+<steptml:parameter name="text">hello world</steptml:parameter>
+</steptml:invoke><｜tool_call_end｜>
+<｜tool_calls_end｜>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["count"], 100);
+    assert_eq!(args["rate"], 2.5);
+    assert_eq!(args["active"], true);
+    assert_eq!(args["optional"], serde_json::Value::Null);
+    assert_eq!(args["text"], "hello world");
+}
+
+#[tokio::test]
+async fn test_step3_streaming() {
+    let mut parser = Step3Parser::new();
+
+    let tools = create_test_tools();
+
+    // Simulate streaming chunks
+    let chunks = vec![
+        "<｜tool_calls_begin｜>\n",
+        "<｜tool_call_begin｜>function",
+        "<｜tool_sep｜><steptml:invoke name=\"calc\">",
+        "\n<steptml:parameter name=\"x\">10</steptml:parameter>",
+        "\n<steptml:parameter name=\"y\">20</steptml:parameter>",
+        "\n</steptml:invoke><｜tool_call_end｜>",
+        "\n<｜tool_calls_end｜>",
+    ];
+
+    let mut found_complete = false;
+
+    for chunk in chunks {
+        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
+
+        if !result.calls.is_empty() {
+            if let Some(name) = &result.calls[0].name {
+                assert_eq!(name, "calc");
+                found_complete = true;
+            }
+        }
+    }
+
+    assert!(found_complete);
+}
+
+#[test]
+fn test_step3_format_detection() {
+    let parser = Step3Parser::new();
+
+    // Should detect Step3 format
+    assert!(parser.has_tool_markers("<｜tool_calls_begin｜>"));
+    assert!(parser.has_tool_markers("text with <｜tool_calls_begin｜> marker"));
+
+    // Should not detect other formats
+    assert!(!parser.has_tool_markers("[TOOL_CALLS]"));
+    assert!(!parser.has_tool_markers("<tool_call>"));
+    assert!(!parser.has_tool_markers("plain text"));
+}
+
+#[tokio::test]
+async fn test_step3_nested_steptml() {
+    let parser = Step3Parser::new();
+
+    let input = r#"<｜tool_calls_begin｜>
+<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="config">
+<steptml:parameter name="settings">{"nested": {"key": "value"}}</steptml:parameter>
+<steptml:parameter name="array">[1, 2, 3]</steptml:parameter>
+</steptml:invoke><｜tool_call_end｜>
+<｜tool_calls_end｜>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "config");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert!(args["settings"].is_object());
+    assert!(args["array"].is_array());
+}
+
+#[tokio::test]
+async fn test_step3_python_literals() {
+    let parser = Step3Parser::new();
+
+    let input = r#"<｜tool_calls_begin｜>
+<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="test">
+<steptml:parameter name="bool_true">True</steptml:parameter>
+<steptml:parameter name="bool_false">False</steptml:parameter>
+<steptml:parameter name="none_value">None</steptml:parameter>
+</steptml:invoke><｜tool_call_end｜>
+<｜tool_calls_end｜>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["bool_true"], true);
+    assert_eq!(args["bool_false"], false);
+    assert_eq!(args["none_value"], serde_json::Value::Null);
+}
+
+#[tokio::test]
+async fn test_steptml_format() {
+    let parser = Step3Parser::new();
+
+    let input = r#"Text before.
+<｜tool_calls_begin｜>
+<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="search">
+<steptml:parameter name="query">rust lang</steptml:parameter>
+<steptml:parameter name="limit">10</steptml:parameter>
+</steptml:invoke><｜tool_call_end｜>
+<｜tool_calls_end｜>Text after."#;
+
+    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(normal_text, "Text before.\n");
+    assert_eq!(tools[0].function.name, "search");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["query"], "rust lang");
+    assert_eq!(args["limit"], 10);
+    // TODO: Verify normal text extraction
+}
+
+#[tokio::test]
+async fn test_json_parameter_values() {
+    let parser = Step3Parser::new();
+
+    let input = r#"<｜tool_calls_begin｜>
+<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="config">
+<steptml:parameter name="settings">{"nested": {"value": true}}</steptml:parameter>
+<steptml:parameter name="items">[1, 2, 3]</steptml:parameter>
+</steptml:invoke><｜tool_call_end｜>
+<｜tool_calls_end｜>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert!(args["settings"].is_object());
+    assert!(args["items"].is_array());
+}
+
+#[tokio::test]
+async fn test_step3_parameter_with_angle_brackets() {
+    let parser = Step3Parser::new();
+
+    let input = r#"<｜tool_calls_begin｜>
+<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="compare">
+<steptml:parameter name="expression">a < b && b > c</steptml:parameter>
+<steptml:parameter name="context">comparison test</steptml:parameter>
+</steptml:invoke><｜tool_call_end｜>
+<｜tool_calls_end｜>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0].function.name, "compare");
+
+    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
+    assert_eq!(args["expression"], "a < b && b > c");
+    assert_eq!(args["context"], "comparison test");
+}
+
+#[tokio::test]
+async fn test_step3_empty_function_name() {
+    let parser = Step3Parser::new();
+
+    let input = r#"<｜tool_calls_begin｜>
+<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="">
+<steptml:parameter name="param">value</steptml:parameter>
+</steptml:invoke><｜tool_call_end｜>
+<｜tool_calls_end｜>"#;
+
+    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
+    assert_eq!(tools.len(), 0); // Should reject empty function name
+}

sgl-model-gateway/tests/wasm_test.rs

@@ -0,0 +1,819 @@
+//! 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 sgl_model_gateway::{
+    app_context::AppContext,
+    config::RouterConfig,
+    core::workflow::{
+        create_wasm_module_registration_workflow, create_wasm_module_removal_workflow,
+    },
+    routers::RouterFactory,
+    server::{build_app, AppState},
+    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();
+
+    // Initialize registries
+    use sgl_model_gateway::{
+        core::{LoadMonitor, WorkerRegistry},
+        data_connector::{
+            MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage,
+        },
+        policies::PolicyRegistry,
+    };
+
+    let worker_registry = Arc::new(WorkerRegistry::new());
+    let policy_registry = Arc::new(PolicyRegistry::new(config.policy.clone()));
+
+    // Initialize storage backends
+    let response_storage = Arc::new(MemoryResponseStorage::new());
+    let conversation_storage = Arc::new(MemoryConversationStorage::new());
+    let conversation_item_storage = Arc::new(MemoryConversationItemStorage::new());
+
+    // Initialize load monitor
+    let load_monitor = Some(Arc::new(LoadMonitor::new(
+        worker_registry.clone(),
+        policy_registry.clone(),
+        client.clone(),
+        config.worker_startup_check_interval_secs,
+    )));
+
+    // Create empty OnceLock for worker job queue, workflow engine, and mcp manager
+    use std::sync::OnceLock;
+    let worker_job_queue = Arc::new(OnceLock::new());
+    let workflow_engine = 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(None)
+            .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_engine(workflow_engine)
+            .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 = sgl_model_gateway::core::JobQueue::new(
+        sgl_model_gateway::core::JobQueueConfig::default(),
+        weak_context,
+    );
+    app_context
+        .worker_job_queue
+        .set(job_queue)
+        .expect("JobQueue should only be initialized once");
+
+    // Initialize WorkflowEngine and register workflows
+    use sgl_model_gateway::core::workflow::{
+        create_worker_registration_workflow, create_worker_removal_workflow, WorkflowEngine,
+    };
+    let engine = Arc::new(WorkflowEngine::new());
+    engine.register_workflow(create_worker_registration_workflow(&config));
+    engine.register_workflow(create_worker_removal_workflow());
+    engine.register_workflow(create_wasm_module_registration_workflow());
+    engine.register_workflow(create_wasm_module_removal_workflow());
+    app_context
+        .workflow_engine
+        .set(engine)
+        .expect("WorkflowEngine should only be initialized once");
+
+    // Initialize MCP manager with empty config
+    use sgl_model_gateway::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,
+    });
+
+    let request_id_headers = vec!["x-request-id".to_string(), "x-correlation-id".to_string()];
+
+    let app = build_app(
+        app_state,
+        sgl_model_gateway::middleware::AuthConfig { api_key: None },
+        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(
+                sgl_model_gateway::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(sgl_model_gateway::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(sgl_model_gateway::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(
+                sgl_model_gateway::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(sgl_model_gateway::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(
+                sgl_model_gateway::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(sgl_model_gateway::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(
+                sgl_model_gateway::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(
+                sgl_model_gateway::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(sgl_model_gateway::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(
+                sgl_model_gateway::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(
+                sgl_model_gateway::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(sgl_model_gateway::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 engine = app_context
+        .workflow_engine
+        .get()
+        .expect("Workflow engine should be initialized");
+
+    // Create workflow context for registration
+    use sgl_model_gateway::core::workflow::{
+        steps::WasmModuleConfigRequest, WorkflowContext, WorkflowId, WorkflowInstanceId,
+    };
+
+    let descriptor = WasmModuleDescriptor {
+        name: "test_execution_module".to_string(),
+        file_path: wasm_file_path.clone(),
+        module_type: WasmModuleType::Middleware,
+        attach_points: vec![WasmModuleAttachPoint::Middleware(
+            sgl_model_gateway::wasm::module::MiddlewareAttachPoint::OnRequest,
+        )],
+        add_result: None,
+    };
+
+    let config_request = WasmModuleConfigRequest { descriptor };
+    let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
+    workflow_context.set_arc("wasm_module_config", Arc::new(config_request));
+    workflow_context.set_arc("app_context", app_context.clone());
+
+    // Start workflow
+    let instance_id = engine
+        .start_workflow(
+            WorkflowId::new("wasm_module_registration"),
+            workflow_context,
+        )
+        .await
+        .expect("Failed to start workflow");
+
+    // Wait for workflow to complete
+    let timeout = Duration::from_secs(30);
+    let start = std::time::Instant::now();
+    let mut module_uuid: Option<Uuid> = None;
+
+    loop {
+        if start.elapsed() > timeout {
+            panic!("Workflow timeout");
+        }
+
+        let state = engine
+            .get_status(instance_id)
+            .expect("Failed to get workflow status");
+
+        match state.status {
+            sgl_model_gateway::core::workflow::WorkflowStatus::Completed => {
+                // Extract module UUID from context
+                if let Some(uuid_arc) = state.context.get::<Uuid>("module_uuid") {
+                    module_uuid = Some(*uuid_arc.as_ref());
+                }
+                break;
+            }
+            sgl_model_gateway::core::workflow::WorkflowStatus::Failed => {
+                panic!("Workflow failed: {:?}", state);
+            }
+            _ => {
+                tokio::time::sleep(Duration::from_millis(100)).await;
+            }
+        }
+    }
+
+    let module_uuid = module_uuid.expect("Module UUID should be in context");
+
+    // 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 sgl_model_gateway::wasm::{
+        spec::sgl::router::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(
+        sgl_model_gateway::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"
+    );
+}

sgl-model-gateway/tests/workflow_test.rs

@@ -0,0 +1,320 @@
+//! Integration tests for workflow engine
+
+use std::{
+    sync::{
+        atomic::{AtomicU32, Ordering},
+        Arc,
+    },
+    time::Duration,
+};
+
+use sgl_model_gateway::core::workflow::*;
+use tokio::time::sleep;
+
+// Test step that counts invocations
+struct CountingStep {
+    counter: Arc<AtomicU32>,
+    should_succeed_after: u32,
+}
+
+#[async_trait::async_trait]
+impl StepExecutor for CountingStep {
+    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+        let count = self.counter.fetch_add(1, Ordering::SeqCst) + 1;
+
+        // Store count in context
+        context.set("execution_count", count);
+
+        if count >= self.should_succeed_after {
+            Ok(StepResult::Success)
+        } else {
+            Err(WorkflowError::StepFailed {
+                step_id: StepId::new("counting_step"),
+                message: format!("Not ready yet, attempt {}", count),
+            })
+        }
+    }
+}
+
+// Test step that always succeeds
+struct AlwaysSucceedStep;
+
+#[async_trait::async_trait]
+impl StepExecutor for AlwaysSucceedStep {
+    async fn execute(&self, _context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+        Ok(StepResult::Success)
+    }
+}
+
+#[tokio::test]
+async fn test_simple_workflow_execution() {
+    let engine = WorkflowEngine::new();
+
+    // Subscribe to events for logging
+    engine
+        .event_bus()
+        .subscribe(Arc::new(LoggingSubscriber))
+        .await;
+
+    // Create a simple workflow
+    let workflow = WorkflowDefinition::new("test_workflow", "Simple Test Workflow")
+        .add_step(StepDefinition::new(
+            "step1",
+            "First Step",
+            Arc::new(AlwaysSucceedStep),
+        ))
+        .add_step(StepDefinition::new(
+            "step2",
+            "Second Step",
+            Arc::new(AlwaysSucceedStep),
+        ));
+
+    let workflow_id = workflow.id.clone();
+    engine.register_workflow(workflow);
+
+    // Start workflow
+    let instance_id = engine
+        .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
+        .await
+        .unwrap();
+
+    // Wait for completion
+    sleep(Duration::from_millis(100)).await;
+
+    // Check status
+    let state = engine.get_status(instance_id).unwrap();
+    assert_eq!(state.status, WorkflowStatus::Completed);
+    assert_eq!(state.step_states.len(), 2);
+}
+
+#[tokio::test]
+async fn test_workflow_with_retry() {
+    let engine = WorkflowEngine::new();
+    engine
+        .event_bus()
+        .subscribe(Arc::new(LoggingSubscriber))
+        .await;
+
+    let counter = Arc::new(AtomicU32::new(0));
+
+    // Create workflow with retry logic
+    let workflow = WorkflowDefinition::new("retry_workflow", "Workflow with Retry").add_step(
+        StepDefinition::new(
+            "retry_step",
+            "Step that retries",
+            Arc::new(CountingStep {
+                counter: Arc::clone(&counter),
+                should_succeed_after: 3,
+            }),
+        )
+        .with_retry(RetryPolicy {
+            max_attempts: 5,
+            backoff: BackoffStrategy::Fixed(Duration::from_millis(10)),
+        })
+        .with_timeout(Duration::from_secs(5)),
+    );
+
+    let workflow_id = workflow.id.clone();
+    engine.register_workflow(workflow);
+
+    // Start workflow
+    let instance_id = engine
+        .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
+        .await
+        .unwrap();
+
+    // Wait for completion
+    sleep(Duration::from_millis(500)).await;
+
+    // Check that step was retried and eventually succeeded
+    let state = engine.get_status(instance_id).unwrap();
+    assert_eq!(state.status, WorkflowStatus::Completed);
+
+    let step_state = state.step_states.get(&StepId::new("retry_step")).unwrap();
+    assert_eq!(step_state.status, StepStatus::Succeeded);
+    assert_eq!(step_state.attempt, 3); // Should have taken 3 attempts
+
+    // Verify counter
+    assert_eq!(counter.load(Ordering::SeqCst), 3);
+}
+
+#[tokio::test]
+async fn test_workflow_failure_after_max_retries() {
+    let engine = WorkflowEngine::new();
+    engine
+        .event_bus()
+        .subscribe(Arc::new(LoggingSubscriber))
+        .await;
+
+    let counter = Arc::new(AtomicU32::new(0));
+
+    // Create workflow that will fail
+    let workflow = WorkflowDefinition::new("failing_workflow", "Workflow that Fails").add_step(
+        StepDefinition::new(
+            "failing_step",
+            "Step that always fails",
+            Arc::new(CountingStep {
+                counter: Arc::clone(&counter),
+                should_succeed_after: 10, // Will never succeed within max_attempts
+            }),
+        )
+        .with_retry(RetryPolicy {
+            max_attempts: 3,
+            backoff: BackoffStrategy::Fixed(Duration::from_millis(10)),
+        })
+        .with_failure_action(FailureAction::FailWorkflow),
+    );
+
+    let workflow_id = workflow.id.clone();
+    engine.register_workflow(workflow);
+
+    // Start workflow
+    let instance_id = engine
+        .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
+        .await
+        .unwrap();
+
+    // Wait for completion
+    sleep(Duration::from_millis(500)).await;
+
+    // Check that workflow failed
+    let state = engine.get_status(instance_id).unwrap();
+    assert_eq!(state.status, WorkflowStatus::Failed);
+
+    let step_state = state.step_states.get(&StepId::new("failing_step")).unwrap();
+    assert_eq!(step_state.status, StepStatus::Failed);
+    assert_eq!(step_state.attempt, 3); // Should have tried 3 times
+
+    // Verify counter
+    assert_eq!(counter.load(Ordering::SeqCst), 3);
+}
+
+#[tokio::test]
+async fn test_workflow_continue_on_failure() {
+    let engine = WorkflowEngine::new();
+    engine
+        .event_bus()
+        .subscribe(Arc::new(LoggingSubscriber))
+        .await;
+
+    let counter = Arc::new(AtomicU32::new(0));
+
+    // Create workflow where first step fails but workflow continues
+    let workflow = WorkflowDefinition::new("continue_workflow", "Continue on Failure")
+        .add_step(
+            StepDefinition::new(
+                "failing_step",
+                "Step that fails",
+                Arc::new(CountingStep {
+                    counter: Arc::clone(&counter),
+                    should_succeed_after: 10,
+                }),
+            )
+            .with_retry(RetryPolicy {
+                max_attempts: 2,
+                backoff: BackoffStrategy::Fixed(Duration::from_millis(10)),
+            })
+            .with_failure_action(FailureAction::ContinueNextStep),
+        )
+        .add_step(StepDefinition::new(
+            "success_step",
+            "Step that succeeds",
+            Arc::new(AlwaysSucceedStep),
+        ));
+
+    let workflow_id = workflow.id.clone();
+    engine.register_workflow(workflow);
+
+    // Start workflow
+    let instance_id = engine
+        .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
+        .await
+        .unwrap();
+
+    // Wait for completion
+    sleep(Duration::from_millis(500)).await;
+
+    // Workflow should complete despite first step failing
+    let state = engine.get_status(instance_id).unwrap();
+    assert_eq!(state.status, WorkflowStatus::Completed);
+
+    // First step should be skipped
+    let step1_state = state.step_states.get(&StepId::new("failing_step")).unwrap();
+    assert_eq!(step1_state.status, StepStatus::Skipped);
+
+    // Second step should succeed
+    let step2_state = state.step_states.get(&StepId::new("success_step")).unwrap();
+    assert_eq!(step2_state.status, StepStatus::Succeeded);
+}
+
+#[tokio::test]
+async fn test_workflow_context_sharing() {
+    let engine = WorkflowEngine::new();
+
+    struct ContextWriterStep {
+        key: String,
+        value: String,
+    }
+
+    #[async_trait::async_trait]
+    impl StepExecutor for ContextWriterStep {
+        async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+            context.set(self.key.clone(), self.value.clone());
+            Ok(StepResult::Success)
+        }
+    }
+
+    struct ContextReaderStep {
+        key: String,
+        expected_value: String,
+    }
+
+    #[async_trait::async_trait]
+    impl StepExecutor for ContextReaderStep {
+        async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+            let value: Arc<String> = context
+                .get(&self.key)
+                .ok_or_else(|| WorkflowError::ContextValueNotFound(self.key.clone()))?;
+
+            if *value == self.expected_value {
+                Ok(StepResult::Success)
+            } else {
+                Err(WorkflowError::StepFailed {
+                    step_id: StepId::new("reader"),
+                    message: format!("Expected {}, got {}", self.expected_value, value),
+                })
+            }
+        }
+    }
+
+    let workflow = WorkflowDefinition::new("context_workflow", "Context Sharing Test")
+        .add_step(StepDefinition::new(
+            "writer",
+            "Write to context",
+            Arc::new(ContextWriterStep {
+                key: "test_key".to_string(),
+                value: "test_value".to_string(),
+            }),
+        ))
+        .add_step(StepDefinition::new(
+            "reader",
+            "Read from context",
+            Arc::new(ContextReaderStep {
+                key: "test_key".to_string(),
+                expected_value: "test_value".to_string(),
+            }),
+        ));
+
+    let workflow_id = workflow.id.clone();
+    engine.register_workflow(workflow);
+
+    let instance_id = engine
+        .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
+        .await
+        .unwrap();
+
+    sleep(Duration::from_millis(100)).await;
+
+    let state = engine.get_status(instance_id).unwrap();
+    assert_eq!(state.status, WorkflowStatus::Completed);
+}