[Bugfix][Rust Frontend] Select earliest-completing stop string (#50200)
Signed-off-by: Samuel Laferriere <9342524+samlaf@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com>
Samuel Laferriere committed
Jul 30, 2026 at 20:33 UTC
d6938b7408ccc9e34b57b684c04f3c103e23bda6
1 file changed
+50
-3
rust/src/text/src/output/decoded.rs
+50
-3
@@ -306,6 +306,12 @@ pub async fn decoded_text_event_stream(
306
/// If stop string matches, returns tuple
307
/// (index into stop string vec, byte index of first byte of stop string in
308
/// output)
309
+///
310
+/// When several stop strings match within the newly generated text (for example
311
+/// when a step appends multiple tokens, or one token decodes to several
312
+/// characters), the stop string that completes earliest in the text is selected,
313
+/// so the result matches appending one character at a time. Ties are broken by
314
+/// stop-list order.
315
fn matches_stop_string(stops: &[String], output: &str, new_bytes: usize) -> Option<(usize, usize)> {
316
// We compare byte subslices to avoid utf8 boundary problem
317
let output = output.as_bytes();
@@ -314,12 +320,15 @@ fn matches_stop_string(stops: &[String], output: &str, new_bytes: usize) -> Opti
320
.iter()
321
.map(|ss| (ss.as_bytes(), ss.len(), next_off.saturating_sub(ss.len())))
322
.enumerate()
317
- .find_map(|(ss_idx, (ss, len, start_off))| {
323
+ .filter_map(|(ss_idx, (ss, len, start_off))| {
324
output[start_off..]
325
.windows(len)
326
.position(|w| w == ss)
321
- .map(|pos| (ss_idx, start_off + pos))
327
+ .map(|pos| (ss_idx, start_off + pos, start_off + pos + len))
328
})
329
+ // `min_by_key` keeps the first minimum, so ties fall back to stop-list order.
330
+ .min_by_key(|&(_, _, end)| end)
331
+ .map(|(ss_idx, start, _)| (ss_idx, start))
332
}
333
334
#[cfg(test)]
@@ -524,11 +533,49 @@ mod tests {
533
#[test]
534
fn stop_string_matches_first_of_multiple() {
535
let stops = vec!["wor".to_string(), "say".to_string()];
527
- // "say" appears earlier but "wor" is checked first (index 0)
536
+ // Only "wor" can complete in the 1-new-byte window; the "say" at index 0
537
+ // is behind the window start and is not a candidate.
538
let result = matches_stop_string(&stops, "say wor", 1);
539
assert_eq!(result, Some((0, 4)));
540
}
541
542
+ /// Several stop strings can land in the same window when a step appends
543
+ /// multiple tokens (speculative decoding) or when one token decodes to
544
+ /// several characters. The earliest-completing one must win over stop-list
545
+ /// order, so that a batched step agrees with character-at-a-time appending.
546
+ #[test]
547
+ fn stop_string_earliest_completing_wins_regardless_of_list_order() {
548
+ // " The user is a": " is a" (5 bytes) was appended in one step. Both
549
+ // "is" (index 10, completes at 12) and " a" (index 12, completes at 14)
550
+ // land in the window. "is" completes earlier, so it wins either order.
551
+ for stops in [vec!["a", "is"], vec!["is", "a"]] {
552
+ let owned: Vec<String> = stops.iter().map(|s| s.to_string()).collect();
553
+ let is_idx = stops.iter().position(|s| *s == "is").unwrap();
554
+ let result = matches_stop_string(&owned, " The user is a", " is a".len());
555
+ assert_eq!(result, Some((is_idx, 10)), "stop list order {stops:?}");
556
+ }
557
+ }
558
+
559
+ /// Both stops complete at the same offset, so stop-list order decides. This
560
+ /// pins `min_by_key`'s first-minimum behavior, which is what implements the
561
+ /// tie-break.
562
+ #[test]
563
+ fn stop_string_ties_broken_by_list_order() {
564
+ for (stops, expected) in [(vec!["ab", "b"], (0, 0)), (vec!["b", "ab"], (0, 1))] {
565
+ let owned: Vec<String> = stops.iter().map(|s| s.to_string()).collect();
566
+ let result = matches_stop_string(&owned, "ab", "ab".len());
567
+ assert_eq!(result, Some(expected), "stop list order {stops:?}");
568
+ }
569
+ }
570
+
571
+ #[test]
572
+ fn stop_string_completion_position_not_start_position() {
573
+ // "b" starts later than "abc" but completes earlier, so it must win.
574
+ let stops = vec!["abc".to_string(), "b".to_string()];
575
+ let result = matches_stop_string(&stops, "abc", 3);
576
+ assert_eq!(result, Some((1, 1)));
577
+ }
578
+
579
#[test]
580
fn stop_string_matches_second_of_multiple() {
581
let stops = vec!["xyz".to_string(), "wor".to_string()];