[router][grpc] Support tool call parser in streaming (#11160)

sgl-router/src/tool_parser/tests.rs

@@ -5,64 +5,27 @@ use crate::tool_parser::partial_json::{
 };
 use crate::tool_parser::traits::ToolParser;
 
-#[test]
-fn test_parse_state_new() {
-    let state = ParseState::new();
-    assert_eq!(state.phase, ParsePhase::Searching);
-    assert_eq!(state.buffer, "");
-    assert_eq!(state.consumed, 0);
-    assert_eq!(state.bracket_depth, 0);
-    assert!(!state.in_string);
-    assert!(!state.escape_next);
+#[tokio::test]
+async fn test_tool_parser_factory() {
+    let factory = ToolParserFactory::new();
+
+    // Test that we can get a pooled parser
+    let pooled_parser = factory.get_pooled("gpt-4");
+    let parser = pooled_parser.lock().await;
+    assert!(parser.detect_format(r#"{"name": "test", "arguments": {}}"#));
 }
 
-#[test]
-fn test_parse_state_process_char() {
-    let mut state = ParseState::new();
+#[tokio::test]
+async fn test_tool_parser_factory_model_mapping() {
+    let factory = ToolParserFactory::new();
 
-    state.process_char('{');
-    assert_eq!(state.bracket_depth, 1);
+    // Test model mapping
+    factory.registry().map_model("test-model", "json");
 
-    state.process_char('}');
-    assert_eq!(state.bracket_depth, 0);
-
-    state.process_char('"');
-    assert!(state.in_string);
-
-    state.process_char('"');
-    assert!(!state.in_string);
-
-    state.process_char('"');
-    state.process_char('\\');
-    assert!(state.escape_next);
-
-    state.process_char('"');
-    assert!(!state.escape_next);
-    assert!(state.in_string); // Still in string because quote was escaped
-}
-
-#[test]
-fn test_parser_registry() {
-    let registry = ParserRegistry::new();
-
-    assert!(!registry.list_mappings().is_empty());
-
-    let mappings = registry.list_mappings();
-    let has_gpt = mappings.iter().any(|(m, _)| m.starts_with("gpt"));
-    assert!(has_gpt);
-}
-
-#[test]
-fn test_parser_registry_pattern_matching() {
-    let mut registry = ParserRegistry::new_for_testing();
-
-    registry.map_model("test-model", "json");
-
-    let mappings = registry.list_mappings();
-    let has_test = mappings
-        .iter()
-        .any(|(m, p)| *m == "test-model" && *p == "json");
-    assert!(has_test);
+    // Get parser for the test model
+    let pooled_parser = factory.get_pooled("test-model");
+    let parser = pooled_parser.lock().await;
+    assert!(parser.detect_format(r#"{"name": "test", "arguments": {}}"#));
 }
 
 #[test]
@@ -165,37 +128,7 @@ fn test_compute_diff() {
     assert_eq!(compute_diff("test", "hello"), "hello");
 }
 
-#[test]
-fn test_stream_result_variants() {
-    let result = StreamResult::Incomplete;
-    matches!(result, StreamResult::Incomplete);
-
-    let result = StreamResult::ToolName {
-        index: 0,
-        name: "test".to_string(),
-    };
-    if let StreamResult::ToolName { index, name } = result {
-        assert_eq!(index, 0);
-        assert_eq!(name, "test");
-    } else {
-        panic!("Expected ToolName variant");
-    }
-
-    let tool = ToolCall {
-        id: "123".to_string(),
-        r#type: "function".to_string(),
-        function: FunctionCall {
-            name: "test".to_string(),
-            arguments: "{}".to_string(),
-        },
-    };
-    let result = StreamResult::ToolComplete(tool.clone());
-    if let StreamResult::ToolComplete(t) = result {
-        assert_eq!(t.id, "123");
-    } else {
-        panic!("Expected ToolComplete variant");
-    }
-}
+// NOTE: test_stream_result_variants removed - StreamResult enum replaced by StreamingParseResult
 
 #[test]
 fn test_partial_tool_call() {
@@ -310,14 +243,12 @@ fn test_json_parser_format_detection() {
 }
 
 #[tokio::test]
-async fn test_registry_with_json_parser() {
-    let registry = ParserRegistry::new();
-
-    // JSON parser should be registered by default
-    assert!(registry.has_parser("json"));
+async fn test_factory_with_json_parser() {
+    let factory = ToolParserFactory::new();
 
     // Should get JSON parser for OpenAI models
-    let parser = registry.get_parser("gpt-4-turbo").unwrap();
+    let pooled_parser = factory.get_pooled("gpt-4-turbo");
+    let parser = pooled_parser.lock().await;
 
     let input = r#"{"name": "test", "arguments": {"x": 1}}"#;
     let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
@@ -546,62 +477,6 @@ mod edge_cases {
         assert!(tools[0].function.arguments.contains("null"));
     }
 
-    #[tokio::test]
-    async fn test_streaming_with_partial_chunks() {
-        let parser = JsonParser::new();
-
-        let mut state1 = ParseState::new();
-        let partial = r#"{"#;
-        let result = parser
-            .parse_incremental(partial, &mut state1)
-            .await
-            .unwrap();
-        assert!(
-            matches!(result, StreamResult::Incomplete),
-            "Should return Incomplete for just opening brace"
-        );
-
-        let mut state2 = ParseState::new();
-        let complete = r#"{"name": "get_weather", "arguments": {"location": "SF"}}"#;
-        let result = parser
-            .parse_incremental(complete, &mut state2)
-            .await
-            .unwrap();
-
-        match result {
-            StreamResult::ToolComplete(tool) => {
-                assert_eq!(tool.function.name, "get_weather");
-                let args: serde_json::Value =
-                    serde_json::from_str(&tool.function.arguments).unwrap();
-                assert_eq!(args["location"], "SF");
-            }
-            _ => panic!("Expected ToolComplete for complete JSON"),
-        }
-
-        // The PartialJson parser can complete partial JSON by filling in missing values
-        let mut state3 = ParseState::new();
-        let partial_with_name = r#"{"name": "test", "argum"#;
-        let result = parser
-            .parse_incremental(partial_with_name, &mut state3)
-            .await
-            .unwrap();
-
-        match result {
-            StreamResult::ToolComplete(tool) => {
-                assert_eq!(tool.function.name, "test");
-                // Arguments will be empty object since "argum" is incomplete
-                assert_eq!(tool.function.arguments, "{}");
-            }
-            StreamResult::ToolName { name, .. } => {
-                assert_eq!(name, "test");
-            }
-            StreamResult::Incomplete => {
-                // Also acceptable if parser decides to wait
-            }
-            _ => panic!("Unexpected result for partial JSON with name"),
-        }
-    }
-
     #[tokio::test]
     async fn test_special_json_values() {
         let parser = JsonParser::new();