aft/bash_background/
watches.rs

1use std::collections::{HashMap, HashSet};
2use std::io::{Read, Seek, SeekFrom};
3
4use super::persistence::ValidatedArtifact;
5
6use regex::{Regex, RegexBuilder};
7use serde::Serialize;
8
9const MAX_WATCHES_PER_TASK: usize = 8;
10const CONTEXT_BEFORE: usize = 100;
11const CONTEXT_AFTER: usize = 500;
12const SCAN_OVERLAP_BYTES: usize = 8 * 1024;
13
14pub const WATCH_TARGET_ERASED_TEXT: &str = "watch target erased";
15pub const WATCH_TARGET_ERASED_CONTEXT: &str =
16    "watch target erased: the background task row was erased before the watch reached a normal terminal result";
17pub const WATCH_TASK_EXIT_TEXT: &str = "watch task exited";
18
19#[derive(Debug, Clone)]
20pub struct WatchSpec {
21    pub watch_id: String,
22    pub task_id: String,
23    pub pattern: WatchPattern,
24    pub once: bool,
25}
26
27#[derive(Debug, Clone)]
28pub enum WatchPattern {
29    Substring(String),
30    Regex(Regex),
31}
32
33impl WatchPattern {
34    pub fn regex(pattern: &str) -> Result<Self, regex::Error> {
35        RegexBuilder::new(pattern)
36            .multi_line(true)
37            .build()
38            .map(Self::Regex)
39    }
40
41    pub fn kind_name(&self) -> &'static str {
42        match self {
43            Self::Substring(_) => "substring",
44            Self::Regex(_) => "regex",
45        }
46    }
47
48    pub fn pattern_text(&self) -> &str {
49        match self {
50            Self::Substring(text) => text,
51            Self::Regex(regex) => regex.as_str(),
52        }
53    }
54
55    pub fn from_persisted(kind: &str, pattern: &str) -> Result<Self, String> {
56        match kind {
57            "substring" => Ok(Self::Substring(pattern.to_string())),
58            "regex" => Self::regex(pattern).map_err(|error| error.to_string()),
59            other => Err(format!("unknown pattern kind {other}")),
60        }
61    }
62}
63
64#[derive(Debug, Clone, Serialize)]
65pub struct PatternMatch {
66    pub watch_id: String,
67    pub task_id: String,
68    pub match_text: String,
69    pub match_offset: u64,
70    pub context: String,
71    pub once: bool,
72}
73
74#[derive(Debug, Default)]
75pub struct WatchRegistry {
76    watches: HashMap<String, Vec<WatchSpec>>,
77    scan_cursors: HashMap<String, u64>,
78    scan_overlaps: HashMap<String, Vec<u8>>,
79    controlled_tasks: HashSet<String>,
80    matched_tasks: HashSet<String>,
81    /// Process-local tombstones keep `bash_status` informative after the durable
82    /// task and watch rows are gone. They intentionally disappear on restart.
83    erased_notifications: HashSet<String>,
84    next_watch: u64,
85}
86
87impl WatchRegistry {
88    pub fn register(
89        &mut self,
90        task_id: String,
91        pattern: WatchPattern,
92        once: bool,
93    ) -> Result<String, &'static str> {
94        let watches = self.watches.entry(task_id.clone()).or_default();
95        if watches.len() >= MAX_WATCHES_PER_TASK {
96            return Err("too_many_watches");
97        }
98        self.controlled_tasks.insert(task_id.clone());
99        self.next_watch = self.next_watch.wrapping_add(1);
100        let watch_id = format!("watch-{:08x}", self.next_watch);
101        watches.push(WatchSpec {
102            watch_id: watch_id.clone(),
103            task_id,
104            pattern,
105            once,
106        });
107        Ok(watch_id)
108    }
109
110    /// Re-arm a previously persisted watch after bridge/daemon restart.
111    ///
112    /// `scanning=false` means a once-watch already matched and is only kept so
113    /// a pending notification can be re-delivered until ack; it does not scan.
114    pub fn restore(
115        &mut self,
116        watch_id: String,
117        task_id: String,
118        pattern: WatchPattern,
119        once: bool,
120        scanning: bool,
121    ) -> Result<(), &'static str> {
122        self.note_restored_watch_id(&watch_id);
123        self.controlled_tasks.insert(task_id.clone());
124        if !scanning {
125            self.matched_tasks.insert(task_id);
126            return Ok(());
127        }
128        let watches = self.watches.entry(task_id.clone()).or_default();
129        if watches.iter().any(|watch| watch.watch_id == watch_id) {
130            return Ok(());
131        }
132        if watches.len() >= MAX_WATCHES_PER_TASK {
133            return Err("too_many_watches");
134        }
135        watches.push(WatchSpec {
136            watch_id,
137            task_id,
138            pattern,
139            once,
140        });
141        Ok(())
142    }
143
144    fn note_restored_watch_id(&mut self, watch_id: &str) {
145        let Some(hex) = watch_id.strip_prefix("watch-") else {
146            return;
147        };
148        if let Ok(value) = u64::from_str_radix(hex, 16) {
149            if value >= self.next_watch {
150                self.next_watch = value;
151            }
152        }
153    }
154
155    pub fn unregister(&mut self, task_id: &str, watch_id: &str) {
156        if let Some(watches) = self.watches.get_mut(task_id) {
157            watches.retain(|watch| watch.watch_id != watch_id);
158            if watches.is_empty() {
159                self.watches.remove(task_id);
160            }
161        }
162    }
163
164    pub fn file_cursor(&self, cursor_key: &str) -> Option<u64> {
165        self.scan_cursors.get(cursor_key).copied()
166    }
167
168    pub fn watch_ids(&self, task_id: &str) -> Vec<String> {
169        self.watches
170            .get(task_id)
171            .map(|watches| watches.iter().map(|watch| watch.watch_id.clone()).collect())
172            .unwrap_or_default()
173    }
174
175    pub fn watched_task_ids(&self) -> Vec<String> {
176        self.watches.keys().cloned().collect()
177    }
178
179    /// Reconcile process-local watch state with the durable rows shared by actors.
180    pub fn reconcile_watch_ids(
181        &mut self,
182        task_id: &str,
183        keep: &HashSet<String>,
184        has_pending_match: bool,
185    ) {
186        let watches_empty = if let Some(watches) = self.watches.get_mut(task_id) {
187            watches.retain(|watch| keep.contains(&watch.watch_id));
188            watches.is_empty()
189        } else {
190            true
191        };
192        if watches_empty {
193            self.clear_task(task_id);
194        } else if has_pending_match {
195            self.matched_tasks.insert(task_id.to_string());
196        } else {
197            self.matched_tasks.remove(task_id);
198        }
199    }
200
201    pub fn clear_task(&mut self, task_id: &str) {
202        self.watches.remove(task_id);
203        self.controlled_tasks.remove(task_id);
204        self.matched_tasks.remove(task_id);
205        let prefix = format!("{task_id}:");
206        self.scan_cursors
207            .retain(|key, _| key != task_id && !key.starts_with(&prefix));
208        self.scan_overlaps
209            .retain(|key, _| key != task_id && !key.starts_with(&prefix));
210    }
211
212    pub fn terminalize_erased_task(&mut self, task_id: &str) -> Vec<String> {
213        let watch_ids = self.watch_ids(task_id);
214        self.clear_task(task_id);
215        watch_ids
216            .into_iter()
217            .filter(|watch_id| {
218                self.erased_notifications
219                    .insert(format!("{task_id}:{watch_id}"))
220            })
221            .collect()
222    }
223
224    pub fn has_erased_task(&self, task_id: &str) -> bool {
225        let prefix = format!("{task_id}:");
226        self.erased_notifications
227            .iter()
228            .any(|notification| notification.starts_with(&prefix))
229    }
230
231    pub fn forget_erased_task(&mut self, task_id: &str) {
232        let prefix = format!("{task_id}:");
233        self.erased_notifications
234            .retain(|notification| !notification.starts_with(&prefix));
235    }
236
237    pub fn has_controlled_task(&self, task_id: &str) -> bool {
238        self.controlled_tasks.contains(task_id)
239    }
240
241    pub fn has_matched_task(&self, task_id: &str) -> bool {
242        self.matched_tasks.contains(task_id)
243    }
244
245    pub fn active_count(&self, task_id: &str) -> usize {
246        self.watches.get(task_id).map_or(0, Vec::len)
247    }
248
249    pub fn prime_file_cursor(&mut self, cursor_key: &str, file: &ValidatedArtifact) {
250        if self.scan_cursors.contains_key(cursor_key) {
251            return;
252        }
253        let len = file.len().unwrap_or(0);
254        self.scan_cursors.insert(cursor_key.to_string(), len);
255    }
256
257    pub fn set_file_cursor(&mut self, cursor_key: &str, offset: u64) {
258        self.scan_cursors.insert(cursor_key.to_string(), offset);
259        self.scan_overlaps.remove(cursor_key);
260    }
261
262    pub fn scan_file_new_bytes(
263        &mut self,
264        cursor_key: &str,
265        task_id: &str,
266        file: &mut ValidatedArtifact,
267    ) -> Vec<PatternMatch> {
268        if self.active_count(task_id) == 0 {
269            return Vec::new();
270        }
271        let cursor = self
272            .scan_cursors
273            .get(cursor_key)
274            .copied()
275            .unwrap_or_else(|| {
276                // Start at current EOF so a newly registered watch does not match old spill content.
277                file.len().unwrap_or(0)
278            });
279        if file.seek(SeekFrom::Start(cursor)).is_err() {
280            return Vec::new();
281        }
282        let mut bytes = Vec::new();
283        if file.read_to_end(&mut bytes).is_err() || bytes.is_empty() {
284            self.scan_cursors.insert(cursor_key.to_string(), cursor);
285            return Vec::new();
286        }
287        let next = cursor.saturating_add(bytes.len() as u64);
288        self.scan_cursors.insert(cursor_key.to_string(), next);
289        self.scan_new_bytes_at(cursor_key, task_id, &bytes, cursor)
290    }
291
292    pub fn scan_new_bytes(&mut self, task_id: &str, bytes: &[u8]) -> Vec<PatternMatch> {
293        let base = self.scan_cursors.get(task_id).copied().unwrap_or(0);
294        self.scan_cursors
295            .insert(task_id.to_string(), base.saturating_add(bytes.len() as u64));
296        self.scan_new_bytes_at(task_id, task_id, bytes, base)
297    }
298
299    fn scan_new_bytes_at(
300        &mut self,
301        cursor_key: &str,
302        task_id: &str,
303        bytes: &[u8],
304        base_offset: u64,
305    ) -> Vec<PatternMatch> {
306        let Some(watches) = self.watches.get(task_id).cloned() else {
307            return Vec::new();
308        };
309        let overlap = self
310            .scan_overlaps
311            .get(cursor_key)
312            .cloned()
313            .unwrap_or_default();
314        let prefix_len = overlap.len();
315        let mut scan_bytes = Vec::with_capacity(prefix_len.saturating_add(bytes.len()));
316        scan_bytes.extend_from_slice(&overlap);
317        scan_bytes.extend_from_slice(bytes);
318        let text = String::from_utf8_lossy(&scan_bytes);
319        let scan_base_offset = base_offset.saturating_sub(prefix_len as u64);
320        let mut matches = Vec::new();
321        let mut remove_once = Vec::new();
322        for watch in watches {
323            if let Some((text_start, text_end, source_start, matched)) =
324                find_match(&watch.pattern, &text, &scan_bytes, prefix_len)
325            {
326                self.matched_tasks.insert(task_id.to_string());
327                matches.push(PatternMatch {
328                    watch_id: watch.watch_id.clone(),
329                    task_id: watch.task_id.clone(),
330                    match_text: matched,
331                    match_offset: scan_base_offset.saturating_add(source_start as u64),
332                    context: context_snippet(&text, text_start, text_end),
333                    once: watch.once,
334                });
335                if watch.once {
336                    remove_once.push(watch.watch_id);
337                }
338            }
339        }
340        for watch_id in remove_once {
341            self.unregister(task_id, &watch_id);
342        }
343        let keep = scan_bytes.len().min(SCAN_OVERLAP_BYTES);
344        self.scan_overlaps.insert(
345            cursor_key.to_string(),
346            scan_bytes[scan_bytes.len().saturating_sub(keep)..].to_vec(),
347        );
348        matches
349    }
350}
351
352fn find_match(
353    pattern: &WatchPattern,
354    text: &str,
355    source: &[u8],
356    min_source_end_exclusive: usize,
357) -> Option<(usize, usize, usize, String)> {
358    let make_match = |start: usize, end: usize, matched: String| {
359        let source_end = source_offset_for_lossy_utf8(source, end);
360        (source_end > min_source_end_exclusive).then(|| {
361            (
362                start,
363                end,
364                source_offset_for_lossy_utf8(source, start),
365                matched,
366            )
367        })
368    };
369    match pattern {
370        WatchPattern::Substring(needle) => {
371            if needle.is_empty() {
372                return None;
373            }
374            text.match_indices(needle).find_map(|(start, matched)| {
375                make_match(start, start + matched.len(), needle.clone())
376            })
377        }
378        WatchPattern::Regex(regex) => regex.find_iter(text).find_map(|matched| {
379            make_match(matched.start(), matched.end(), matched.as_str().to_string())
380        }),
381    }
382}
383
384fn source_offset_for_lossy_utf8(source: &[u8], lossy_offset: usize) -> usize {
385    let mut source_start = 0;
386    let mut rendered_start = 0;
387    while source_start < source.len() {
388        match std::str::from_utf8(&source[source_start..]) {
389            Ok(valid) => {
390                return source_start
391                    .saturating_add(lossy_offset.saturating_sub(rendered_start).min(valid.len()));
392            }
393            Err(error) => {
394                let valid_len = error.valid_up_to();
395                let rendered_valid_end = rendered_start.saturating_add(valid_len);
396                if lossy_offset <= rendered_valid_end {
397                    return source_start
398                        .saturating_add(lossy_offset.saturating_sub(rendered_start));
399                }
400                source_start = source_start.saturating_add(valid_len);
401                rendered_start = rendered_valid_end;
402
403                let invalid_len = error
404                    .error_len()
405                    .unwrap_or_else(|| source.len().saturating_sub(source_start));
406                if lossy_offset < rendered_start.saturating_add('\u{FFFD}'.len_utf8()) {
407                    return source_start;
408                }
409                source_start = source_start.saturating_add(invalid_len);
410                rendered_start = rendered_start.saturating_add('\u{FFFD}'.len_utf8());
411            }
412        }
413    }
414    source.len()
415}
416
417fn context_snippet(text: &str, start: usize, end: usize) -> String {
418    let before_start = text[..start]
419        .char_indices()
420        .rev()
421        .nth(CONTEXT_BEFORE)
422        .map(|(idx, _)| idx)
423        .unwrap_or(0);
424    let after_end = text[end..]
425        .char_indices()
426        .nth(CONTEXT_AFTER)
427        .map(|(idx, _)| end + idx)
428        .unwrap_or(text.len());
429    text[before_start..after_end].replace('\r', "")
430}
431
432#[cfg(test)]
433mod tests {
434    use super::*;
435
436    #[test]
437    fn once_watch_self_removes_after_match() {
438        let mut registry = WatchRegistry::default();
439        let task_id = "bash-1".to_string();
440        registry
441            .register(
442                task_id.clone(),
443                WatchPattern::Substring("READY".into()),
444                true,
445            )
446            .unwrap();
447        assert_eq!(registry.scan_new_bytes(&task_id, b"READY\n").len(), 1);
448        assert_eq!(registry.active_count(&task_id), 0);
449    }
450
451    #[test]
452    fn sticky_watch_fires_multiple_times() {
453        let mut registry = WatchRegistry::default();
454        let task_id = "bash-1".to_string();
455        registry
456            .register(
457                task_id.clone(),
458                WatchPattern::Substring("READY".into()),
459                false,
460            )
461            .unwrap();
462        assert_eq!(registry.scan_new_bytes(&task_id, b"READY\n").len(), 1);
463        assert_eq!(registry.scan_new_bytes(&task_id, b"READY\n").len(), 1);
464        assert_eq!(registry.active_count(&task_id), 1);
465    }
466
467    #[test]
468    fn cap_8_watches_per_task_rejects_9th() {
469        let mut registry = WatchRegistry::default();
470        for _ in 0..8 {
471            registry
472                .register("bash-1".into(), WatchPattern::Substring("x".into()), true)
473                .unwrap();
474        }
475        assert_eq!(
476            registry.register("bash-1".into(), WatchPattern::Substring("x".into()), true),
477            Err("too_many_watches")
478        );
479    }
480
481    #[test]
482    fn regex_pattern_matches_with_capture() {
483        let mut registry = WatchRegistry::default();
484        let task_id = "bash-1".to_string();
485        registry
486            .register(
487                task_id.clone(),
488                WatchPattern::regex("port (\\d+)").unwrap(),
489                true,
490            )
491            .unwrap();
492        let hits = registry.scan_new_bytes(&task_id, b"listening on port 3000\n");
493        assert_eq!(hits[0].match_text, "port 3000");
494    }
495
496    #[test]
497    fn substring_pattern_can_span_scans() {
498        let mut registry = WatchRegistry::default();
499        let task_id = "bash-1".to_string();
500        registry
501            .register(
502                task_id.clone(),
503                WatchPattern::Substring("READY".into()),
504                true,
505            )
506            .unwrap();
507
508        assert!(registry.scan_new_bytes(&task_id, b"RE").is_empty());
509        let hits = registry.scan_new_bytes(&task_id, b"ADY\n");
510
511        assert_eq!(hits.len(), 1);
512        assert_eq!(hits[0].match_text, "READY");
513        assert_eq!(hits[0].match_offset, 0);
514    }
515
516    #[test]
517    fn regex_pattern_can_span_scans() {
518        let mut registry = WatchRegistry::default();
519        let task_id = "bash-1".to_string();
520        registry
521            .register(
522                task_id.clone(),
523                WatchPattern::regex("ready: \\d{4}").unwrap(),
524                true,
525            )
526            .unwrap();
527
528        assert!(registry
529            .scan_new_bytes(&task_id, b"prefix ready: 4")
530            .is_empty());
531        let hits = registry.scan_new_bytes(&task_id, b"242\n");
532
533        assert_eq!(hits.len(), 1);
534        assert_eq!(hits[0].match_text, "ready: 4242");
535        assert_eq!(hits[0].match_offset, 7);
536    }
537
538    #[test]
539    fn overlap_does_not_repeat_fully_previous_match() {
540        let mut registry = WatchRegistry::default();
541        let task_id = "bash-1".to_string();
542        registry
543            .register(
544                task_id.clone(),
545                WatchPattern::Substring("READY".into()),
546                false,
547            )
548            .unwrap();
549
550        assert_eq!(registry.scan_new_bytes(&task_id, b"READY").len(), 1);
551        assert!(registry.scan_new_bytes(&task_id, b"\n").is_empty());
552    }
553
554    #[test]
555    fn lossy_utf8_indices_are_mapped_back_to_source_offsets() {
556        let mut registry = WatchRegistry::default();
557        let task_id = "bash-1".to_string();
558        registry
559            .register(
560                task_id.clone(),
561                WatchPattern::Substring("READY".into()),
562                false,
563            )
564            .unwrap();
565
566        let hits = registry.scan_new_bytes(&task_id, b"\xffREADY");
567
568        assert_eq!(hits.len(), 1);
569        assert_eq!(hits[0].match_offset, 1);
570        assert!(registry.scan_new_bytes(&task_id, b"\n").is_empty());
571    }
572}
aft/bash_background/watches.rs

aft/bash_background/
watches.rs
