1use std::time::Duration;
14
15use regex::Regex;
16
17use crate::tools::strip_ansi;
18
19use super::config::ShellConfig;
20
21#[derive(Debug, Clone)]
27pub enum ReadinessStrategy {
28 Timeout,
30 Prompt,
32 Hybrid,
34}
35
36impl ReadinessStrategy {
37 #[allow(clippy::should_implement_trait)]
38 pub fn from_str(s: &str) -> Self {
39 match s.to_lowercase().as_str() {
40 "timeout" => ReadinessStrategy::Timeout,
41 "prompt" => ReadinessStrategy::Prompt,
42 _ => ReadinessStrategy::Hybrid,
43 }
44 }
45}
46
47#[derive(Debug, PartialEq)]
49pub enum ReadinessResult {
50 Ready,
52 Waiting,
54 SilenceTimeout,
56 MaxTimeout,
58}
59
60pub struct ReadinessDetector {
66 strategy: ReadinessStrategy,
67 patterns: Vec<Regex>,
68 silence_timeout: Duration,
69 max_timeout: Duration,
70}
71
72const PROMPT_TAIL_LEN: usize = 200;
74
75impl ReadinessDetector {
76 pub fn new(
81 strategy: ReadinessStrategy,
82 patterns_str: &[String],
83 silence_timeout_ms: u64,
84 max_timeout_ms: u64,
85 ) -> Self {
86 let patterns: Vec<Regex> = patterns_str
87 .iter()
88 .filter_map(|p| match Regex::new(p) {
89 Ok(re) => Some(re),
90 Err(e) => {
91 tracing::warn!(pattern = %p, error = %e, "skipping invalid prompt regex");
92 None
93 }
94 })
95 .collect();
96
97 let strategy = if patterns.is_empty() {
99 match strategy {
100 ReadinessStrategy::Prompt | ReadinessStrategy::Hybrid => {
101 tracing::warn!(
102 "no valid prompt patterns — falling back to Timeout strategy"
103 );
104 ReadinessStrategy::Timeout
105 }
106 other => other,
107 }
108 } else {
109 strategy
110 };
111
112 Self {
113 strategy,
114 patterns,
115 silence_timeout: Duration::from_millis(silence_timeout_ms),
116 max_timeout: Duration::from_millis(max_timeout_ms),
117 }
118 }
119
120 pub fn from_config(config: &ShellConfig) -> Self {
122 let strategy = crate::tools::shell::readiness::ReadinessStrategy::from_str(&config.readiness_strategy);
123 Self::new(
124 strategy,
125 &config.prompt_patterns,
126 config.readiness_timeout_ms,
127 config.max_readiness_timeout_ms,
128 )
129 }
130
131 pub fn check(
137 &self,
138 output: &str,
139 silence_elapsed: Duration,
140 total_elapsed: Duration,
141 ) -> ReadinessResult {
142 if total_elapsed >= self.max_timeout {
144 return ReadinessResult::MaxTimeout;
145 }
146
147 match &self.strategy {
148 ReadinessStrategy::Timeout => {
149 if silence_elapsed >= self.silence_timeout {
150 ReadinessResult::SilenceTimeout
151 } else {
152 ReadinessResult::Waiting
153 }
154 }
155 ReadinessStrategy::Prompt => {
156 if self.matches_prompt(output) {
157 ReadinessResult::Ready
158 } else {
159 ReadinessResult::Waiting
160 }
161 }
162 ReadinessStrategy::Hybrid => {
163 if self.matches_prompt(output) {
164 ReadinessResult::Ready
165 } else if silence_elapsed >= self.silence_timeout {
166 ReadinessResult::SilenceTimeout
167 } else {
168 ReadinessResult::Waiting
169 }
170 }
171 }
172 }
173
174 pub fn matches_prompt(&self, output: &str) -> bool {
180 if output.is_empty() || self.patterns.is_empty() {
181 return false;
182 }
183
184 let tail = if output.len() > PROMPT_TAIL_LEN {
190 let mut start = output.len() - PROMPT_TAIL_LEN;
191 while start > 0 && !output.is_char_boundary(start) {
192 start -= 1;
193 }
194 &output[start..]
195 } else {
196 output
197 };
198
199 let clean = strip_ansi(tail);
200
201 let last_line = clean
203 .lines()
204 .rev()
205 .find(|l| !l.trim().is_empty())
206 .unwrap_or("");
207
208 if last_line.is_empty() {
209 return false;
210 }
211
212 self.patterns.iter().any(|re| re.is_match(last_line))
213 }
214}
215
216#[cfg(test)]
221mod tests {
222 use super::*;
223
224 fn timeout_detector(silence_ms: u64) -> ReadinessDetector {
226 ReadinessDetector::new(ReadinessStrategy::Timeout, &[], silence_ms, 10_000)
227 }
228
229 fn prompt_detector() -> ReadinessDetector {
231 let config = ShellConfig::default();
232 ReadinessDetector::new(
233 ReadinessStrategy::Prompt,
234 &config.prompt_patterns,
235 config.readiness_timeout_ms,
236 config.max_readiness_timeout_ms,
237 )
238 }
239
240 fn hybrid_detector() -> ReadinessDetector {
242 ReadinessDetector::from_config(&ShellConfig::default())
243 }
244
245 #[test]
247 fn timeout_strategy_silence_triggers() {
248 let det = timeout_detector(300);
249 let result = det.check(
250 "some output\n",
251 Duration::from_millis(301),
252 Duration::from_millis(500),
253 );
254 assert_eq!(result, ReadinessResult::SilenceTimeout);
255 }
256
257 #[test]
259 fn timeout_strategy_still_waiting() {
260 let det = timeout_detector(300);
261 let result = det.check(
262 "some output\n",
263 Duration::from_millis(100),
264 Duration::from_millis(200),
265 );
266 assert_eq!(result, ReadinessResult::Waiting);
267 }
268
269 #[test]
271 fn prompt_strategy_dollar_ready() {
272 let det = prompt_detector();
273 let result = det.check(
274 "user@host:~$ ",
275 Duration::from_millis(0),
276 Duration::from_millis(50),
277 );
278 assert_eq!(result, ReadinessResult::Ready);
279 }
280
281 #[test]
283 fn prompt_strategy_python_ready() {
284 let det = prompt_detector();
285 let result = det.check(
286 "Python 3.11.0\n>>> ",
287 Duration::from_millis(0),
288 Duration::from_millis(50),
289 );
290 assert_eq!(result, ReadinessResult::Ready);
291 }
292
293 #[test]
295 fn prompt_strategy_no_match_silence_fallback() {
296 let det = prompt_detector();
297 let result = det.check(
299 "compiling crate...\n",
300 Duration::from_millis(500),
301 Duration::from_millis(1000),
302 );
303 assert_eq!(result, ReadinessResult::Waiting);
304 }
305
306 #[test]
308 fn prompt_strategy_no_match_waiting() {
309 let det = prompt_detector();
310 let result = det.check(
311 "compiling crate...\n",
312 Duration::from_millis(100),
313 Duration::from_millis(200),
314 );
315 assert_eq!(result, ReadinessResult::Waiting);
316 }
317
318 #[test]
320 fn hybrid_prompt_match_before_silence() {
321 let det = hybrid_detector();
322 let result = det.check(
323 "welcome\nuser@host:~$ ",
324 Duration::from_millis(10), Duration::from_millis(50),
326 );
327 assert_eq!(result, ReadinessResult::Ready);
328 }
329
330 #[test]
332 fn hybrid_silence_fallback() {
333 let det = hybrid_detector();
334 let result = det.check(
335 "running long task...\n",
336 Duration::from_millis(500),
337 Duration::from_millis(1000),
338 );
339 assert_eq!(result, ReadinessResult::SilenceTimeout);
340 }
341
342 #[test]
344 fn max_timeout_always_wins() {
345 for det in [timeout_detector(300), prompt_detector(), hybrid_detector()] {
346 let result = det.check(
347 "user@host:~$ ",
348 Duration::from_millis(0),
349 Duration::from_millis(10_001),
350 );
351 assert_eq!(result, ReadinessResult::MaxTimeout);
352 }
353 }
354
355 #[test]
357 fn matches_prompt_common_patterns() {
358 let det = hybrid_detector();
359 let prompts = [
360 "user@host:~$ ",
361 "root@server:/var# ",
362 ">>> ",
363 "(gdb) ",
364 "Password: ",
365 ];
366 for prompt in &prompts {
367 assert!(
368 det.matches_prompt(prompt),
369 "expected match for prompt: {:?}",
370 prompt,
371 );
372 }
373 }
374
375 #[test]
377 fn invalid_regex_skipped_no_panic() {
378 let patterns = vec![
379 "[invalid(".into(), r"[$#] $".into(), ];
382 let det = ReadinessDetector::new(
383 ReadinessStrategy::Hybrid,
384 &patterns,
385 300,
386 10_000,
387 );
388 assert!(det.matches_prompt("user@host:~$ "));
390 }
391
392 #[test]
394 fn all_invalid_patterns_fallback_to_timeout() {
395 let patterns = vec!["[broken(".into(), "(also[bad".into()];
396 let det = ReadinessDetector::new(
397 ReadinessStrategy::Hybrid,
398 &patterns,
399 300,
400 10_000,
401 );
402 assert!(!det.matches_prompt("user@host:~$ "));
404 let result = det.check(
405 "user@host:~$ ",
406 Duration::from_millis(301),
407 Duration::from_millis(500),
408 );
409 assert_eq!(result, ReadinessResult::SilenceTimeout);
410 }
411
412 #[test]
414 fn empty_output_no_match() {
415 let det = hybrid_detector();
416 assert!(!det.matches_prompt(""));
417 }
418
419 #[test]
421 fn ansi_stripped_before_matching() {
422 let det = hybrid_detector();
423 let ansi_prompt = "\x1b[32muser@host\x1b[0m:\x1b[34m~\x1b[0m$ ";
425 assert!(det.matches_prompt(ansi_prompt));
426 }
427
428 #[test]
430 fn only_last_line_checked() {
431 let det = hybrid_detector();
432 assert!(!det.matches_prompt("user@host:~$ \nstill running..."));
434 assert!(det.matches_prompt("still running...\nuser@host:~$ "));
436 }
437
438 #[test]
443 fn matches_prompt_handles_multibyte_glyph_at_tail_boundary() {
444 let det = hybrid_detector();
445
446 let prefix_len = PROMPT_TAIL_LEN - 1;
450 let mut output = "x".repeat(prefix_len);
451 output.push('●'); output.push_str("\n~/repo on main\nuser@host:~$ ");
453
454 assert!(det.matches_prompt(&output));
456 }
457
458 #[test]
459 fn matches_prompt_handles_4byte_emoji_at_tail_boundary() {
460 let det = hybrid_detector();
461
462 for offset in 0..4 {
464 let prefix_len = PROMPT_TAIL_LEN - offset;
465 let mut output = "a".repeat(prefix_len);
466 output.push('🔥'); output.push_str("\nuser@host:~$ ");
468 assert!(det.matches_prompt(&output), "panicked at offset {}", offset);
469 }
470 }
471}