refactor: New config types for tool calls (#4575) (#4857)

Signed-off-by: ayushag <[email protected]>
Co-authored-by: William Zhang <[email protected]>

Author: ayushag-nv

lib/llm/src/protocols/openai/chat_completions/jail.rs

@@ -887,14 +887,14 @@ impl JailedStreamBuilder {
             if let Some(config) = parser_map.get(parser_name.as_str()) {
                 // Auto-populate start sequences if none configured
                 if self.jail_start_sequences.is_empty() {
-                    self.jail_start_sequences = config.json.tool_call_start_tokens.clone();
+                    self.jail_start_sequences = config.parser_config.tool_call_start_tokens();
                 }
 
                 // Auto-populate end sequences if none configured
                 if self.jail_end_sequences.is_empty() {
                     self.jail_end_sequences = config
-                        .json
-                        .tool_call_end_tokens
+                        .parser_config
+                        .tool_call_end_tokens()
                         .iter()
                         .filter(|&s| !s.is_empty())
                         .cloned()
@@ -914,7 +914,7 @@ impl JailedStreamBuilder {
             let parser_map = get_tool_parser_map();
             if let Some(config) = parser_map.get(parser_name.as_str()) {
                 // Add start tokens from the parser config
-                all_patterns.extend(config.json.tool_call_start_tokens.clone());
+                all_patterns.extend(config.parser_config.tool_call_start_tokens());
             }
         }
 

lib/llm/tests/test_jail.rs

@@ -1935,6 +1935,239 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn test_jailed_stream_qwen3_coder_parser() {
+        // Input:
+        // "I'll call a function. "
+        // + "<tool_call><function=get_weather><parameter=location>San Francisco</parameter><parameter=unit>celsius</parameter></function></tool_call>"
+        // + " Done."
+        // Expected output: 3 chunks [Content(), ToolCall(), Content()]
+        let chunks = vec![
+            create_mock_response_chunk("I'll call a function. ".to_string(), 0),
+            create_mock_response_chunk("<tool_call>".to_string(), 0),
+            create_mock_response_chunk("<function=get_weather>".to_string(), 0),
+            create_mock_response_chunk(
+                "<parameter=location>San Francisco</parameter>".to_string(),
+                0,
+            ),
+            create_mock_response_chunk("<parameter=unit>celsius</parameter>".to_string(), 0),
+            create_mock_response_chunk("</function>".to_string(), 0),
+            create_mock_response_chunk("</tool_call>".to_string(), 0),
+            create_mock_response_chunk(" Done.".to_string(), 0),
+        ];
+
+        let input_stream = stream::iter(chunks);
+
+        let jail = JailedStream::builder()
+            .tool_call_parser("qwen3_coder")
+            .build();
+
+        let results: Vec<_> = jail.apply_with_finish_reason(input_stream).collect().await;
+
+        assert_eq!(
+            results.len(),
+            3,
+            "Should have content, tool call, and trailing content"
+        );
+
+        // Verify exact output structure: [Content(), ToolCall(), Content()].
+        test_utils::assert_content(&results[0], "I'll call a function. ");
+        test_utils::assert_tool_call(
+            &results[1],
+            "get_weather",
+            serde_json::json!({"location": "San Francisco", "unit": "celsius"}),
+        );
+        test_utils::assert_content(&results[2], " Done.");
+
+        // Verify content reconstruction excludes tool calls.
+        let reconstructed = test_utils::reconstruct_content(&results);
+        assert_eq!(reconstructed, "I'll call a function.  Done.");
+    }
+
+    #[tokio::test]
+    async fn test_jailed_stream_qwen3_coder_multiple_params() {
+        let chunks = vec![
+            create_mock_response_chunk("Let me search for that. ".to_string(), 0),
+            create_mock_response_chunk(
+                "<tool_call><function=web_search><parameter=query>Rust programming</parameter><parameter=max_results>10</parameter><parameter=filter>recent</parameter></function></tool_call>".to_string(),
+                0,
+            ),
+            create_mock_response_chunk(" Searching now.".to_string(), 0),
+        ];
+
+        let input_stream = stream::iter(chunks);
+        let jail = JailedStream::builder()
+            .tool_call_parser("qwen3_coder")
+            .build();
+
+        let results: Vec<_> = jail.apply_with_finish_reason(input_stream).collect().await;
+
+        assert_eq!(results.len(), 3, "Should have 3 chunks");
+
+        test_utils::assert_content(&results[0], "Let me search for that. ");
+        test_utils::assert_tool_call(
+            &results[1],
+            "web_search",
+            serde_json::json!({
+                "query": "Rust programming",
+                "max_results": 10,
+                "filter": "recent"
+            }),
+        );
+        test_utils::assert_content(&results[2], " Searching now.");
+    }
+
+    #[tokio::test]
+    async fn test_jailed_stream_xml_parser_config_tokens_auto_population() {
+        // Tests that parser config tokens are auto-populated when using `.tool_call_parser()`.
+        // This verifies the jail system reads `tool_call_start_token` and `tool_call_end_token`
+        // from the `qwen3_coder` parser config.
+        let chunks = vec![
+            create_mock_response_chunk("Before tool call. ".to_string(), 0),
+            create_mock_response_chunk("<tool_call>".to_string(), 0), // Default qwen3_coder token
+            create_mock_response_chunk("<function=get_weather>".to_string(), 0),
+            create_mock_response_chunk("<parameter=city>Seattle</parameter>".to_string(), 0),
+            create_mock_response_chunk("</function>".to_string(), 0),
+            create_mock_response_chunk("</tool_call>".to_string(), 0), // Default qwen3_coder token
+            create_mock_response_chunk(" After tool call.".to_string(), 0),
+        ];
+
+        let input_stream = stream::iter(chunks);
+
+        // Create JailedStream using ONLY `.tool_call_parser()`.
+        // This should auto-populate jail sequences from the qwen3_coder config
+        let jail = JailedStream::builder()
+            .tool_call_parser("qwen3_coder")
+            .build();
+
+        let results: Vec<_> = jail.apply_with_finish_reason(input_stream).collect().await;
+
+        assert_eq!(
+            results.len(),
+            3,
+            "Should have content, tool call, and trailing content"
+        );
+
+        test_utils::assert_content(&results[0], "Before tool call. ");
+        test_utils::assert_tool_call(
+            &results[1],
+            "get_weather",
+            serde_json::json!({"city": "Seattle"}),
+        );
+        test_utils::assert_content(&results[2], " After tool call.");
+
+        let reconstructed = test_utils::reconstruct_content(&results);
+        assert_eq!(reconstructed, "Before tool call.  After tool call.");
+    }
+
+    #[tokio::test]
+    async fn test_jailed_stream_xml_manual_sequences_prevent_auto_population() {
+        // Tests that manually setting jail sequences prevents auto-population.
+        // This verifies the builder respects manual configuration over auto-population.
+        //
+        // When custom sequences are set, the default parser tokens (<tool_call>) should
+        // NOT trigger jailing and should pass through as regular content.
+        let chunks = vec![
+            create_mock_response_chunk("Text with ".to_string(), 0),
+            // Default qwen3_coder token - should NOT trigger jailing.
+            create_mock_response_chunk("<tool_call>".to_string(), 0),
+            create_mock_response_chunk("should not jail".to_string(), 0),
+            create_mock_response_chunk("</tool_call>".to_string(), 0),
+            create_mock_response_chunk(" because custom ".to_string(), 0),
+            // Custom marker - this SHOULD trigger jailing since we register it below.
+            create_mock_response_chunk("[[START]]".to_string(), 0),
+            create_mock_response_chunk("jailed content".to_string(), 0),
+            create_mock_response_chunk("[[END]]".to_string(), 0),
+            create_mock_response_chunk(" text.".to_string(), 0),
+        ];
+
+        let input_stream = stream::iter(chunks);
+
+        // Set custom jail sequences - this should prevent auto-population.
+        // The default <tool_call> tokens should NOT trigger jailing.
+        let jail = JailedStream::builder()
+            .jail_start_sequence("[[START]]")
+            .jail_end_sequence("[[END]]")
+            .tool_call_parser("qwen3_coder")
+            .build();
+
+        let results: Vec<_> = jail.apply_with_finish_reason(input_stream).collect().await;
+
+        // The exact number of chunks depends on emission mode (packed vs single-choice-per-chunk)
+        // but we can verify the key behaviors:
+        // 1. Default <tool_call> tokens pass through as content (not jailed)
+        // 2. Custom [[START]]/[[END]] markers trigger jailing
+        // 3. No tool calls are extracted (because jailed content isn't valid XML)
+
+        // Find chunk(s) containing the default tokens that passed through.
+        let default_token_chunks: Vec<_> = results
+            .iter()
+            .filter_map(|r| {
+                r.data
+                    .as_ref()
+                    .and_then(|d| d.choices.first())
+                    .and_then(|c| c.delta.content.as_ref())
+            })
+            .filter(|content| {
+                content.contains("<tool_call>") || content.contains("should not jail")
+            })
+            .collect();
+
+        assert!(
+            !default_token_chunks.is_empty(),
+            "Default <tool_call> should pass through as content when manual sequences are set"
+        );
+
+        // Find chunk containing the jailed content that was released.
+        let jailed_chunk = results
+            .iter()
+            .filter_map(|r| {
+                r.data
+                    .as_ref()
+                    .and_then(|d| d.choices.first())
+                    .and_then(|c| c.delta.content.as_ref())
+            })
+            .find(|content| content.contains("[[START]]") && content.contains("jailed content"));
+
+        assert!(
+            jailed_chunk.is_some(),
+            "Custom markers should trigger jailing and accumulated content should be released"
+        );
+
+        // Since the custom markers include non-XML content, the parser should not extract tool calls.
+        // The accumulated content "[[START]]jailed content[[END]]", although compatible with the
+        // way we configured `jail` above, is not consistent with what `qwen_coder` expects, and
+        // there is (at time of writing) no way to pass a parser instance - only a string that
+        // internally gets mapped to default way of instantiating a particular parser.
+        let tool_call_count = results
+            .iter()
+            .filter(|r| {
+                r.data
+                    .as_ref()
+                    .and_then(|d| d.choices.first())
+                    .and_then(|c| c.delta.tool_calls.as_ref())
+                    .map(|tc| !tc.is_empty())
+                    .unwrap_or(false)
+            })
+            .count();
+
+        assert_eq!(
+            tool_call_count, 0,
+            "Should have 0 tool calls because jailed content doesn't match XML format"
+        );
+
+        // Verify content reconstruction - all original content should be preserved.
+        let reconstructed = test_utils::reconstruct_content(&results);
+        assert!(
+            reconstructed.contains("<tool_call>") && reconstructed.contains("should not jail"),
+            "Reconstructed content should include default tokens that passed through"
+        );
+        assert!(
+            reconstructed.contains("[[START]]") && reconstructed.contains("jailed content"),
+            "Reconstructed content should include jailed content with custom markers"
+        );
+    }
+
     #[tokio::test]
     async fn test_jailed_stream_mistral_false_positive_curly() {
         // Curly brace in normal text should not trigger tool call detection for mistral