[Rust Frontend] Recover buffered text from incomplete tool calls at EOS (#47289)
Signed-off-by: reidliu41 <reid201711@gmail.com>
Reid committed
Jul 3, 2026 at 11:45 UTC
9b8e76589dc8cd3a5c6928e8ca6f1bd5feea0a2b
1 file changed
+61
-9
rust/src/chat/src/output/default/unified.rs
+61
-9
@@ -112,7 +112,8 @@ impl UnifiedParserState {
112
);
113
self.parser_failed = true;
114
self.open_call_index = None;
115
- // TODO: should we reset and emit the buffered text?
115
+ let recovered = self.parser.reset();
116
+ push_text_delta(&mut events, AssistantBlockKind::Text, recovered);
117
}
118
}
119
@@ -268,11 +269,13 @@ pub(crate) async fn unified_event_stream(
269
#[cfg(test)]
270
mod tests {
271
use std::collections::VecDeque;
272
+ use std::sync::Arc;
273
274
use futures::{StreamExt as _, stream};
275
use vllm_parser::reasoning::ReasoningError;
274
- use vllm_parser::tool::ToolCallDelta;
275
- use vllm_parser::unified::{UnifiedParserError, UnifiedParserOutput};
276
+ use vllm_parser::tool::{Tool, ToolCallDelta};
277
+ use vllm_parser::unified::{Gemma4UnifiedParser, UnifiedParserError, UnifiedParserOutput};
278
+ use vllm_tokenizer::test_utils::TestTokenizer;
279
280
use super::unified_event_stream;
281
use crate::event::AssistantBlockKind;
@@ -606,7 +609,7 @@ mod tests {
609
}
610
611
#[tokio::test]
609
- async fn unified_stream_finish_error_closes_parser_without_reset_text() {
612
+ async fn unified_stream_finish_error_recovers_buffered_text() {
613
let events = collect(
614
ScriptedParser::new([ScriptedStep::Output(UnifiedParserOutput::default())])
615
.with_finish_error("buffered"),
@@ -616,11 +619,60 @@ mod tests {
619
620
assert_eq!(
621
events,
619
- vec![AssistantEvent::Done {
620
- usage: vllm_llm::TokenUsage::default(),
621
- finish_reason: crate::FinishReason::Stop(None),
622
- kv_transfer_params: None,
623
- }]
622
+ vec![
623
+ AssistantEvent::TextDelta {
624
+ kind: AssistantBlockKind::Text,
625
+ delta: "buffered".to_string(),
626
+ },
627
+ AssistantEvent::Done {
628
+ usage: vllm_llm::TokenUsage::default(),
629
+ finish_reason: crate::FinishReason::Stop(None),
630
+ kv_transfer_params: None,
631
+ },
632
+ ]
633
+ );
634
+ }
635
+
636
+ #[tokio::test]
637
+ async fn unified_stream_recovers_incomplete_gemma4_tool_call_at_eos() {
638
+ let tokenizer = TestTokenizer::new()
639
+ .with_special_token("<|channel>", 256)
640
+ .with_special_token("<channel|>", 257);
641
+ let tools = vec![Tool {
642
+ name: "write_file".to_string(),
643
+ description: None,
644
+ parameters: serde_json::json!({ "type": "object" }),
645
+ strict: None,
646
+ }];
647
+ let parser = Gemma4UnifiedParser::new(&tools, Arc::new(tokenizer)).unwrap();
648
+ let events = vec![
649
+ decoded_delta("<|tool_call>"),
650
+ decoded_delta("call:write_file{"),
651
+ decoded_delta("content:<|\"|>hello "),
652
+ finished_delta("world<|\"|>"),
653
+ ];
654
+ let stream = stream::iter(events.into_iter().map(Ok));
655
+ let events = unified_event_stream(stream, Box::new(parser))
656
+ .collect::<Vec<_>>()
657
+ .await
658
+ .into_iter()
659
+ .collect::<crate::Result<Vec<_>>>()
660
+ .unwrap();
661
+
662
+ assert_eq!(
663
+ events,
664
+ vec![
665
+ AssistantEvent::TextDelta {
666
+ kind: AssistantBlockKind::Text,
667
+ delta: "<|tool_call>call:write_file{content:<|\"|>hello world<|\"|>"
668
+ .to_string(),
669
+ },
670
+ AssistantEvent::Done {
671
+ usage: vllm_llm::TokenUsage::default(),
672
+ finish_reason: crate::FinishReason::Stop(None),
673
+ kv_transfer_params: None,
674
+ },
675
+ ]
676
);
677
}
678
}