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

aft/bash_background/
watches.rs
