diff --git a/rust/src/chat/src/output/default/unified.rs b/rust/src/chat/src/output/default/unified.rs index b12924371f..e9ba676c8c 100644 --- a/rust/src/chat/src/output/default/unified.rs +++ b/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("", 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::>() + .await + .into_iter() + .collect::>>() + .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, + }, + ] ); } }