[Rust Frontend] Recover buffered text from incomplete tool calls at EOS (#47289)

Signed-off-by: reidliu41 <reid201711@gmail.com>

rust/src/chat/src/output/default/unified.rs

@@ -112,7 +112,8 @@ impl UnifiedParserState {
                 );
                 self.parser_failed = true;
                 self.open_call_index = None;
-                // TODO: should we reset and emit the buffered text?
+                let recovered = self.parser.reset();
+                push_text_delta(&mut events, AssistantBlockKind::Text, recovered);
             }
         }
 
@@ -268,11 +269,13 @@ pub(crate) async fn unified_event_stream(
 #[cfg(test)]
 mod tests {
     use std::collections::VecDeque;
+    use std::sync::Arc;
 
     use futures::{StreamExt as _, stream};
     use vllm_parser::reasoning::ReasoningError;
-    use vllm_parser::tool::ToolCallDelta;
-    use vllm_parser::unified::{UnifiedParserError, UnifiedParserOutput};
+    use vllm_parser::tool::{Tool, ToolCallDelta};
+    use vllm_parser::unified::{Gemma4UnifiedParser, UnifiedParserError, UnifiedParserOutput};
+    use vllm_tokenizer::test_utils::TestTokenizer;
 
     use super::unified_event_stream;
     use crate::event::AssistantBlockKind;
@@ -606,7 +609,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn unified_stream_finish_error_closes_parser_without_reset_text() {
+    async fn unified_stream_finish_error_recovers_buffered_text() {
         let events = collect(
             ScriptedParser::new([ScriptedStep::Output(UnifiedParserOutput::default())])
                 .with_finish_error("buffered"),
@@ -616,11 +619,60 @@ mod tests {
 
         assert_eq!(
             events,
-            vec![AssistantEvent::Done {
-                usage: vllm_llm::TokenUsage::default(),
-                finish_reason: crate::FinishReason::Stop(None),
-                kv_transfer_params: None,
-            }]
+            vec![
+                AssistantEvent::TextDelta {
+                    kind: AssistantBlockKind::Text,
+                    delta: "buffered".to_string(),
+                },
+                AssistantEvent::Done {
+                    usage: vllm_llm::TokenUsage::default(),
+                    finish_reason: crate::FinishReason::Stop(None),
+                    kv_transfer_params: None,
+                },
+            ]
+        );
+    }
+
+    #[tokio::test]
+    async fn unified_stream_recovers_incomplete_gemma4_tool_call_at_eos() {
+        let tokenizer = TestTokenizer::new()
+            .with_special_token("<|channel>", 256)
+            .with_special_token("<channel|>", 257);
+        let tools = vec![Tool {
+            name: "write_file".to_string(),
+            description: None,
+            parameters: serde_json::json!({ "type": "object" }),
+            strict: None,
+        }];
+        let parser = Gemma4UnifiedParser::new(&tools, Arc::new(tokenizer)).unwrap();
+        let events = vec![
+            decoded_delta("<|tool_call>"),
+            decoded_delta("call:write_file{"),
+            decoded_delta("content:<|\"|>hello "),
+            finished_delta("world<|\"|>"),
+        ];
+        let stream = stream::iter(events.into_iter().map(Ok));
+        let events = unified_event_stream(stream, Box::new(parser))
+            .collect::<Vec<_>>()
+            .await
+            .into_iter()
+            .collect::<crate::Result<Vec<_>>>()
+            .unwrap();
+
+        assert_eq!(
+            events,
+            vec![
+                AssistantEvent::TextDelta {
+                    kind: AssistantBlockKind::Text,
+                    delta: "<|tool_call>call:write_file{content:<|\"|>hello world<|\"|>"
+                        .to_string(),
+                },
+                AssistantEvent::Done {
+                    usage: vllm_llm::TokenUsage::default(),
+                    finish_reason: crate::FinishReason::Stop(None),
+                    kv_transfer_params: None,
+                },
+            ]
         );
     }
 }