From f835279dd53472bf715433be2c8b4bb4d1509bbf Mon Sep 17 00:00:00 2001 From: simpleqt <89645338+simpleqt@users.noreply.github.com> Date: Sun, 20 Sep 2026 20:46:10 +0800 Subject: [PATCH] fix(stream): flush still-buffered token ids at end of stream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit StreamDecoder::decode returns None whenever the decoded text ends with U+FFFD, on the assumption that a multi-byte character is still split across tokens and the next token will complete it. The heuristic cannot distinguish decoded text that legitimately ends with U+FFFD (mojibake bytes decode to a real replacement character), and nothing ever flushed the buffer at end of stream — the processor only finished the state machine. A backend stream ending with such a token therefore lost its final buffered text from both content and completion usage (completion_tokens undercounted the buffered ids). StreamDecoder::finish() now flushes whatever remains at end of stream: a genuinely truncated sequence yields its replacement character, a legitimate one is emitted instead of dropped. The processor feeds the flushed text through the state machine (stop sequences still match) before finishing it. Regression test drives Token chunks ('a', U+FFFD) into a Finish and asserts content 'a\u{FFFD}' with completion_tokens 2; it fails on the previous implementation (content 'a'). --- deepseek-recipe/src/stream/decoder.rs | 21 +++++++ deepseek-recipe/src/stream/processor.rs | 76 +++++++++++++++++++++++++ 2 files changed, 97 insertions(+) diff --git a/deepseek-recipe/src/stream/decoder.rs b/deepseek-recipe/src/stream/decoder.rs index 8c694f6..215294b 100644 --- a/deepseek-recipe/src/stream/decoder.rs +++ b/deepseek-recipe/src/stream/decoder.rs @@ -68,4 +68,25 @@ impl StreamDecoder { self.ids.clear(); Ok(Some((content, std::mem::take(&mut self.pending)))) } + + /// Flush any still-buffered ids at the end of the stream. + /// + /// The incompleteness heuristic cannot distinguish a multi-token character + /// still split across tokens from decoded text that legitimately ends with + /// U+FFFD (mojibake bytes decode to a real replacement character). At + /// end-of-stream there is no further token to disambiguate, so whatever + /// remains is emitted as-is — a genuinely truncated sequence yields its + /// replacement character, a legitimate one is no longer dropped. + pub fn finish(&mut self) -> Option<(String, usize)> { + if self.ids.is_empty() { + return None; + } + let content = self + .decoder + .decode_ids(&self.ids, false) + .ok() + .filter(|content| !content.is_empty())?; + self.ids.clear(); + Some((content, std::mem::take(&mut self.pending))) + } } diff --git a/deepseek-recipe/src/stream/processor.rs b/deepseek-recipe/src/stream/processor.rs index 40fc4f5..6affa18 100644 --- a/deepseek-recipe/src/stream/processor.rs +++ b/deepseek-recipe/src/stream/processor.rs @@ -140,6 +140,19 @@ where yield Ok(out); } } + // 流结束时冲刷仍在缓冲的 token：尾部的替换字符（真实 U+FFFD 文本） + // 或截断序列不应从输出与用量里消失。 + if let Some(decoder) = decoder.as_mut() { + if let Some((content, content_tokens)) = decoder.finish() { + completion_usage.completion_tokens += content_tokens; + let actions = state_machine.feed(&content); + stashed.push(Some(content)); + for out in stashed.apply_actions(actions, &mut generator).await { + yield Ok(out); + } + stop_sequence = stashed.take_stop_sequence(); + } + } let actions = state_machine.finish(); for out in stashed.apply_actions(actions, &mut generator).await { yield Ok(out); @@ -367,3 +380,66 @@ fn escape_json_string(content: String) -> String { _ => String::new(), } } + +#[cfg(test)] +mod decoder_flush_tests { + use super::super::decoder::TokenizerDecoder; + use super::*; + + struct ScriptedDecoder; + + /// id 0 -> "a"；id 1 -> U+FFFD（真实文本里的替换字符，非多字节截断）。 + impl TokenizerDecoder for ScriptedDecoder { + fn decode_ids(&self, ids: &[u32], _skip_special_tokens: bool) -> Result { + Ok(ids + .iter() + .map(|id| match id { + 0 => "a", + _ => "\u{FFFD}", + }) + .collect()) + } + } + + struct Collector; + + impl ChunkGenerator for Collector { + type Chunk = OutputChunk; + + async fn generate(&mut self, chunk: OutputChunk) -> Vec { + vec![chunk] + } + } + + /// 流末尾缓冲的 token 必须冲刷：结尾恰好是真实 U+FFFD 的文本不能 + /// 从输出与 completion 用量里消失（此前启发式把它当作未完成字符而丢弃）。 + #[tokio::test] + async fn trailing_replacement_character_is_flushed_at_end_of_stream() { + let inference = tokio_stream::iter(vec![ + InferenceChunk::Token { token_id: 0 }, + InferenceChunk::Token { token_id: 1 }, + InferenceChunk::Finish { + finish_reason: InferenceFinishReason::Stop, + }, + ]); + let processor = StreamProcessor::new(Collector, ParsingOptions::default()) + .with_tokenizer(ScriptedDecoder); + let mut content = String::new(); + let mut usage = None; + let stream = processor.process(inference); + tokio::pin!(stream); + while let Some(event) = stream.next().await { + match event { + Ok(OutputChunk::Raw { content: text }) => content += &text, + Ok(OutputChunk::Finish { + usage: finish_usage, + .. + }) => usage = Some(finish_usage), + Ok(_) => {} + Err(error) => panic!("stream error: {error:?}"), + } + } + assert_eq!(content, "a\u{FFFD}"); + assert_eq!(usage.expect("finish event").completion_tokens, 2); + } +}