[Bugfix][Rust Frontend] Detokenizer: avoid leaking prompt on zero-generated-token completions (#47707)
Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: xiaguan <751080330@qq.com>
JinYan Su committed
Jul 16, 2026 at 17:11 UTC
cc706b05a5e5fcca9f5ba7dab57dcfd18515e27e
1 file changed
+34
-5
rust/src/tokenizer/src/incremental.rs
+34
-5
@@ -158,17 +158,19 @@ impl<T: Tokenizer + ?Sized> IncrementalDecoder for DecodeStream<'_, T> {
158
}
159
160
fn flush(&mut self, truncate_output_to: Option<usize>) -> Result<(Option<String>, String)> {
161
- if !self.ids.is_empty() {
161
+ // If the prefix was never seeded (no push_token was called), `ids`
162
+ // holds only prompt context — decoding it would re-emit prompt text.
163
+ if self.prefix_seeded && !self.ids.is_empty() {
164
let string = self.tokenizer.decode(&self.ids, self.skip_special_tokens)?;
165
let prefix_len = self.prefix.len();
164
- self.ids.clear();
165
- self.prefix.clear();
166
- self.prefix_index = 0;
167
- self.prefix_seeded = true;
166
// Ensure we split at a utf-8 char boundary.
167
self.cumulative_output
168
.push_str(&string[string.floor_char_boundary(prefix_len)..]);
169
}
170
+ self.ids.clear();
171
+ self.prefix.clear();
172
+ self.prefix_index = 0;
173
+ self.prefix_seeded = true;
174
if let Some(truncate_output_to) = truncate_output_to {
175
self.cumulative_output.truncate(truncate_output_to);
176
}
@@ -489,4 +491,31 @@ mod tests {
491
assert_eq!(full_text, "你好A");
492
assert_eq!(out, "你好A");
493
}
494
+
495
+ #[test]
496
+ fn flush_without_push_token_does_not_leak_prompt() {
497
+ let backend = Utf8Backend;
498
+ let prompt: Vec<u32> = b"The quick brown fox jumps over the lazy dog. "
499
+ .iter()
500
+ .cycle()
501
+ .take(7001)
502
+ .map(|&b| b as u32)
503
+ .collect();
504
+ let mut decoder = backend.create_decode_stream(&prompt, false, 0);
505
+
506
+ let (last_chunk, full_text) = decoder.flush(None).unwrap();
507
+ assert_eq!(last_chunk, None);
508
+ assert_eq!(full_text, "");
509
+ }
510
+
511
+ #[test]
512
+ fn flush_without_push_token_does_not_leak_undecodable_prompt_tail() {
513
+ let backend = Utf8Backend;
514
+ let prompt = vec![0xe4, 0xbd];
515
+ let mut decoder = backend.create_decode_stream(&prompt, false, 0);
516
+
517
+ let (last_chunk, full_text) = decoder.flush(None).unwrap();
518
+ assert_eq!(last_chunk, None);
519
+ assert_eq!(full_text, "");
520
+ }
521
}