1 //! Spots a streamed reply that has degenerated into a repetition loop.
2 //!
3 //! A model that falls into degenerate decoding keeps producing words from a
4 //! tiny vocabulary and never emits a tool call or a stop (issue #123: thousands
5 //! of words of `find: list: check: get: …`). It is not an exact repeat, so the
6 //! check looks at how few distinct words a recent window holds, not at literal
7 //! repeated strings.
8
9 use std::collections::{HashMap, VecDeque};
10
11 /// Words in the sliding window.
12 const WINDOW: usize = 200;
13 /// The window is re-checked after this many new words, so the cost stays flat
14 /// on a long reply.
15 const CHECK_EVERY: usize = 25;
16 /// A full window with at most this many distinct words is degenerate. Ordinary
17 /// prose of 200 words holds roughly 90 to 120 distinct ones; the loop in the
18 /// issue holds about 30 and an exact repeat far fewer, so there is room on both
19 /// sides.
20 const MAX_DISTINCT: usize = 40;
21
22 /// Fed the same visible text the user is shown, in order. Words inside fenced
23 /// code blocks are skipped: code, tables and data are legitimately repetitive.
24 #[derive(Default)]
25 pub struct RepetitionGuard {
26 /// Bytes fed so far; the offset of the next char in the caller's text.
27 offset: usize,
28 window: VecDeque<(String, usize)>,
29 counts: HashMap<String, usize>,
30 since_check: usize,
31 word: String,
32 word_start: usize,
33 /// The first up-to-3 non-blank chars of the current line.
34 line_head: String,
35 fence_line: bool,
36 in_code: bool,
37 tripped_at: Option<usize>,
38 }
39
40 impl RepetitionGuard {
41 pub fn new() -> Self {
42 Self::default()
43 }
44
45 /// Feed the next piece of text. Returns true once the reply has tripped,
46 /// and keeps returning true after that.
47 pub fn push(&mut self, text: &str) -> bool {
48 for (i, c) in text.char_indices() {
49 if self.tripped_at.is_some() {
50 break;
51 }
52 self.feed(self.offset + i, c);
53 }
54 self.offset += text.len();
55 self.tripped_at.is_some()
56 }
57
58 /// Byte offset in the fed text where the offending window began. The text
59 /// before it is the reply worth keeping.
60 pub fn loop_start(&self) -> Option<usize> {
61 self.tripped_at
62 }
63
64 fn feed(&mut self, at: usize, c: char) {
65 if c == '\n' {
66 self.end_word();
67 self.line_head.clear();
68 self.fence_line = false;
69 return;
70 }
71 if self.line_head.chars().count() < 3 && !(c.is_whitespace() && self.line_head.is_empty()) {
72 self.line_head.push(c);
73 if self.line_head == "```" {
74 self.in_code = !self.in_code;
75 self.fence_line = true;
76 }
77 }
78 if self.in_code || self.fence_line || !c.is_alphanumeric() {
79 self.end_word();
80 return;
81 }
82 if self.word.is_empty() {
83 self.word_start = at;
84 }
85 self.word.extend(c.to_lowercase());
86 }
87
88 /// Where the loop began inside the window. The window can hold a few words
89 /// of real reply ahead of it, so walk back from the end while words still
90 /// belong to the vocabulary of the window's second half.
91 fn loop_origin(&self) -> usize {
92 let vocab: std::collections::HashSet<&str> = self
93 .window
94 .iter()
95 .skip(WINDOW / 2)
96 .map(|(word, _)| word.as_str())
97 .collect();
98 let mut origin = self.window.front().map_or(0, |(_, start)| *start);
99 for (word, start) in self.window.iter().rev() {
100 if !vocab.contains(word.as_str()) {
101 break;
102 }
103 origin = *start;
104 }
105 origin
106 }
107
108 fn end_word(&mut self) {
109 if self.word.is_empty() {
110 return;
111 }
112 let word = std::mem::take(&mut self.word);
113 *self.counts.entry(word.clone()).or_default() += 1;
114 self.window.push_back((word, self.word_start));
115 if self.window.len() > WINDOW
116 && let Some((old, _)) = self.window.pop_front()
117 && let Some(count) = self.counts.get_mut(&old)
118 {
119 *count -= 1;
120 if *count == 0 {
121 self.counts.remove(&old);
122 }
123 }
124 self.since_check += 1;
125 if self.window.len() == WINDOW
126 && self.since_check >= CHECK_EVERY
127 && self.counts.len() <= MAX_DISTINCT
128 {
129 self.tripped_at = Some(self.loop_origin());
130 }
131 if self.since_check >= CHECK_EVERY {
132 self.since_check = 0;
133 }
134 }
135 }
136
137 /// The shape of the loop in issue #123: a random walk over about 30 words.
138 #[cfg(test)]
139 pub(crate) fn issue_loop_text() -> String {
140 let vocab = [
141 "find:", "list:", "check:", "get:", "find", "them:", "list", "commits", "with", "that",
142 "key:", "show:", "read:", "search:", "open:", "match:", "run:", "see:", "look:", "take:",
143 "grab:", "pull:", "fetch:", "scan:", "test:", "view:", "load:", "parse:", "walk:", "dump:",
144 ];
145 let mut state = 7u32;
146 let mut out = String::new();
147 for _ in 0..1500 {
148 state = state.wrapping_mul(1_103_515_245).wrapping_add(12_345);
149 out.push_str(vocab[(state >> 16) as usize % vocab.len()]);
150 out.push(' ');
151 }
152 out
153 }
154
155 #[cfg(test)]
156 mod tests {
157 use super::*;
158
159 fn prose() -> String {
160 // Distinct words throughout, like a real paragraph.
161 (0..400)
162 .map(|n| format!("word{n}"))
163 .collect::<Vec<_>>()
164 .join(" ")
165 }
166
167 fn feed_in_chunks(guard: &mut RepetitionGuard, text: &str) -> bool {
168 let mut tripped = false;
169 for piece in text.as_bytes().chunks(7) {
170 tripped = guard.push(&String::from_utf8_lossy(piece));
171 }
172 tripped
173 }
174
175 #[test]
176 fn the_issue_loop_trips() {
177 let mut guard = RepetitionGuard::new();
178 assert!(feed_in_chunks(&mut guard, &issue_loop_text()));
179 }
180
181 #[test]
182 fn a_short_exact_loop_trips_and_reports_where_it_began() {
183 let intro = "Let me check the history first. ";
184 let text = format!("{intro}{}", "a b a b ".repeat(200));
185 let mut guard = RepetitionGuard::new();
186 assert!(guard.push(&text));
187 let start = guard.loop_start().unwrap();
188 assert!(start >= intro.len(), "cut {start} is inside the intro");
189 assert!(text.is_char_boundary(start));
190 }
191
192 #[test]
193 fn ordinary_prose_does_not_trip() {
194 let mut guard = RepetitionGuard::new();
195 assert!(!feed_in_chunks(&mut guard, &prose()));
196 assert_eq!(guard.loop_start(), None);
197 }
198
199 #[test]
200 fn a_fenced_block_of_repeated_lines_does_not_trip() {
201 let text = format!("Here:\n```\n{}```\nDone.\n", "foo bar\n".repeat(400));
202 let mut guard = RepetitionGuard::new();
203 assert!(!feed_in_chunks(&mut guard, &text));
204 }
205
206 #[test]
207 fn a_list_of_similar_bullets_does_not_trip() {
208 let text: String = (0..60)
209 .map(|n| format!("- Update the handler for route{n} so it returns item{n} quickly\n"))
210 .collect();
211 let mut guard = RepetitionGuard::new();
212 assert!(!feed_in_chunks(&mut guard, &text));
213 }
214
215 #[test]
216 fn text_after_a_closed_fence_counts_again() {
217 let text = format!("```\nx\n```\n{}", "a b a b ".repeat(200));
218 let mut guard = RepetitionGuard::new();
219 assert!(guard.push(&text));
220 }
221 }