[router][grpc] Fix streaming bugs: empty tool names, state pollution, and panics (#11373)
This commit is contained in:
Chang Su
@@ -1,73 +1,199 @@
|
||||
//! Streaming Parser Tests
|
||||
//! Realistic Streaming Parser Tests
|
||||
//!
|
||||
//! Tests for incremental/streaming parsing capabilities across all parsers
|
||||
//! Tests incremental parsing with realistic char-level chunks (2-5 chars)
|
||||
//! that simulate how LLM tokens actually arrive.
|
||||
//!
|
||||
//! These tests are designed to catch bugs like `{"name": "` being parsed
|
||||
//! as an empty tool name.
|
||||
|
||||
use sglang_router_rs::tool_parser::{
|
||||
JsonParser, LlamaParser, MistralParser, PythonicParser, QwenParser, ToolParser,
|
||||
};
|
||||
use sglang_router_rs::tool_parser::{JsonParser, LlamaParser, QwenParser, ToolParser};
|
||||
|
||||
mod common;
|
||||
use common::create_test_tools;
|
||||
use common::{create_test_tools, streaming_helpers::*};
|
||||
|
||||
// =============================================================================
|
||||
// THE BUG SCENARIO - Most Critical Test
|
||||
// =============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_streaming_simple() {
|
||||
async fn test_json_bug_incomplete_tool_name_string() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let full_json = r#"{"name": "get_weather", "arguments": {"location": "San Francisco"}}"#;
|
||||
|
||||
let result = parser.parse_incremental(full_json, &tools).await.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should have parsed a tool call");
|
||||
assert_eq!(result.calls[0].name, Some("get_weather".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_streaming_array() {
|
||||
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#"{"name": "tool1", "#,
|
||||
r#""arguments": {}}, "#,
|
||||
r#"{"name": "tool2", "#,
|
||||
r#""arguments": {"x": 1"#,
|
||||
r#"}}]"#,
|
||||
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 tool_count = 0;
|
||||
let mut got_tool_name = false;
|
||||
let mut saw_empty_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
for chunk in chunks.iter() {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
for call in result.calls {
|
||||
if call.name.is_some() {
|
||||
tool_count += 1;
|
||||
if let Some(name) = &call.name {
|
||||
if name.is_empty() {
|
||||
saw_empty_name = true;
|
||||
}
|
||||
if name == "search" {
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Current implementation may handle this differently
|
||||
assert!(tool_count <= 2, "Should parse at most 2 tools");
|
||||
assert!(
|
||||
!saw_empty_name,
|
||||
"Parser should NEVER return empty tool name"
|
||||
);
|
||||
assert!(got_tool_name, "Should have parsed tool name correctly");
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// JSON PARSER REALISTIC STREAMING
|
||||
// =============================================================================
|
||||
|
||||
#[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_mistral_streaming() {
|
||||
async fn test_json_strategic_chunks_with_quotes() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let mut parser = MistralParser::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");
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// LLAMA PARSER REALISTIC STREAMING
|
||||
// =============================================================================
|
||||
|
||||
#[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![
|
||||
r#"Here is the result: "#,
|
||||
r#"[TOOL_CALLS] ["#,
|
||||
r#"{"name": "#,
|
||||
r#""search", "#,
|
||||
r#""arguments": "#,
|
||||
r#"{"query": "#,
|
||||
r#""rust lang""#,
|
||||
r#"}}]"#,
|
||||
"<|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;
|
||||
@@ -82,40 +208,47 @@ async fn test_mistral_streaming() {
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have found tool name");
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// QWEN PARSER REALISTIC STREAMING
|
||||
// =============================================================================
|
||||
|
||||
#[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_pythonic_streaming() {
|
||||
async fn test_qwen_xml_tag_arrives_in_parts() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = PythonicParser::new();
|
||||
|
||||
let full_input = r#"[get_weather(city="London", units="celsius")]"#;
|
||||
|
||||
let result = parser.parse_incremental(full_input, &tools).await.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should have parsed a tool call");
|
||||
assert_eq!(result.calls[0].name, Some("get_weather".to_string()));
|
||||
let args: serde_json::Value = serde_json::from_str(&result.calls[0].parameters).unwrap();
|
||||
assert_eq!(args["city"], "London");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_streaming_with_python_tag() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = LlamaParser::new();
|
||||
let mut parser = QwenParser::new();
|
||||
|
||||
let chunks = vec![
|
||||
r#"Let me help. "#,
|
||||
r#"<|python"#,
|
||||
r#"_tag|>"#,
|
||||
r#"{"name": "#,
|
||||
r#""calculate", "#,
|
||||
r#""arguments": "#,
|
||||
r#"{"x": 10}"#,
|
||||
r#"}"#,
|
||||
"<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;
|
||||
@@ -124,191 +257,66 @@ async fn test_llama_streaming_with_python_tag() {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "calculate");
|
||||
assert_eq!(name, "translate");
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have found tool name");
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_streaming() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = QwenParser::new();
|
||||
|
||||
// Note: Parser expects newline after both tags
|
||||
let full_input = "<tool_call>\n{\"name\": \"translate\", \"arguments\": {\"text\": \"hello\", \"to\": \"zh\"}}\n</tool_call>";
|
||||
|
||||
let result = parser.parse_incremental(full_input, &tools).await.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should have parsed a tool call");
|
||||
assert_eq!(result.calls[0].name, Some("translate".to_string()));
|
||||
}
|
||||
// =============================================================================
|
||||
// EDGE CASES WITH REALISTIC CHUNKS
|
||||
// =============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_streaming_incomplete_stays_incomplete() {
|
||||
async fn test_json_very_long_url_in_arguments() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let chunks = vec![r#"{"na"#, r#"me": "#];
|
||||
// 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();
|
||||
assert!(
|
||||
result.calls.is_empty(),
|
||||
"Should return empty calls for partial JSON, got: {:?}",
|
||||
result
|
||||
);
|
||||
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_streaming_buffer_accumulation() {
|
||||
async fn test_json_unicode_arrives_byte_by_byte() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let result1 = parser.parse_incremental(r#"{"na"#, &tools).await.unwrap();
|
||||
let input = r#"{"name": "search", "arguments": {"query": "Hello 世界 🌍"}}"#;
|
||||
let chunks = create_realistic_chunks(input);
|
||||
|
||||
assert!(result1.calls.is_empty(), "Should not parse incomplete JSON");
|
||||
let mut got_tool_name = false;
|
||||
|
||||
let result2 = parser
|
||||
.parse_incremental(r#"me": "test", "arguments": {}}"#, &tools)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(
|
||||
!result2.calls.is_empty(),
|
||||
"Should parse complete JSON after buffering"
|
||||
);
|
||||
assert_eq!(result2.calls[0].name, Some("test".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_streaming_multiple_tools_sequential() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = QwenParser::new();
|
||||
|
||||
let full_input = r#"<tool_call>
|
||||
{"name": "tool1", "arguments": {}}
|
||||
</tool_call>"#;
|
||||
|
||||
let result = parser.parse_incremental(full_input, &tools).await.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should have parsed a tool call");
|
||||
assert_eq!(result.calls[0].name, Some("tool1".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_streaming_reset_after_error() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser1 = JsonParser::new();
|
||||
|
||||
let _ = parser1
|
||||
.parse_incremental(r#"{"name": invalid}"#, &tools)
|
||||
.await;
|
||||
|
||||
// Use a new parser instance for clean state
|
||||
let mut parser2 = JsonParser::new();
|
||||
let result = parser2
|
||||
.parse_incremental(r#"{"name": "test", "arguments": {}}"#, &tools)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should parse valid JSON");
|
||||
assert_eq!(result.calls[0].name, Some("test".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_streaming_with_unicode_chunks() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let full_input = r#"{"name": "translate", "arguments": {"text": "Hello 世界 🌍"}}"#;
|
||||
|
||||
let result = parser.parse_incremental(full_input, &tools).await.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should have parsed a tool call");
|
||||
|
||||
// Check if we got the tool name
|
||||
if let Some(name) = &result.calls[0].name {
|
||||
assert_eq!(name, "translate");
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// In streaming mode, need to make another call to get parameters
|
||||
let result2 = parser.parse_incremental("", &tools).await.unwrap();
|
||||
|
||||
// Parameters should be in either result.calls[1] or result2.calls[0]
|
||||
let params = if result.calls.len() > 1 {
|
||||
&result.calls[1].parameters
|
||||
} else if !result2.calls.is_empty() {
|
||||
&result2.calls[0].parameters
|
||||
} else {
|
||||
&result.calls[0].parameters
|
||||
};
|
||||
|
||||
if !params.is_empty() {
|
||||
let args: serde_json::Value = serde_json::from_str(params).unwrap();
|
||||
assert!(args["text"].as_str().unwrap().contains("世界"));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_streaming_with_partial_chunks() {
|
||||
let mut parser = JsonParser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
let partial = r#"{"#;
|
||||
let result = parser.parse_incremental(partial, &tools).await.unwrap();
|
||||
assert!(
|
||||
result.calls.is_empty(),
|
||||
"Should return empty calls for just opening brace"
|
||||
);
|
||||
|
||||
let mut parser2 = JsonParser::new();
|
||||
let complete = r#"{"name": "get_weather", "arguments": {"location": "SF"}}"#;
|
||||
let result = parser2.parse_incremental(complete, &tools).await.unwrap();
|
||||
|
||||
assert!(
|
||||
!result.calls.is_empty(),
|
||||
"Expected tool call for complete JSON"
|
||||
);
|
||||
assert_eq!(result.calls[0].name.as_ref().unwrap(), "get_weather");
|
||||
|
||||
// In streaming mode, need to make another call to get parameters
|
||||
let result2 = parser2.parse_incremental("", &tools).await.unwrap();
|
||||
|
||||
// Parameters should be in either result.calls[1] or result2.calls[0]
|
||||
let params = if result.calls.len() > 1 {
|
||||
&result.calls[1].parameters
|
||||
} else if !result2.calls.is_empty() {
|
||||
&result2.calls[0].parameters
|
||||
} else {
|
||||
&result.calls[0].parameters
|
||||
};
|
||||
|
||||
if !params.is_empty() {
|
||||
let args: serde_json::Value = serde_json::from_str(params).unwrap();
|
||||
assert_eq!(args["location"], "SF");
|
||||
}
|
||||
|
||||
// The PartialJson parser can complete partial JSON by filling in missing values
|
||||
let mut parser3 = JsonParser::new();
|
||||
let partial_with_name = r#"{"name": "test", "argum"#;
|
||||
let result = parser3
|
||||
.parse_incremental(partial_with_name, &tools)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Parser behavior may vary - either complete with partial data or wait for more
|
||||
if !result.calls.is_empty() {
|
||||
assert_eq!(result.calls[0].name.as_ref().unwrap(), "test");
|
||||
}
|
||||
assert!(got_tool_name, "Should have parsed with unicode");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user