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((start, end, matched)) = find_match(&watch.pattern, &text, prefix_len) {
275                self.matched_tasks.insert(task_id.to_string());
276                matches.push(PatternMatch {
277                    watch_id: watch.watch_id.clone(),
278                    task_id: watch.task_id.clone(),
279                    match_text: matched,
280                    match_offset: scan_base_offset.saturating_add(start as u64),
281                    context: context_snippet(&text, start, end),
282                    once: watch.once,
283                });
284                if watch.once {
285                    remove_once.push(watch.watch_id);
286                }
287            }
288        }
289        for watch_id in remove_once {
290            self.unregister(task_id, &watch_id);
291        }
292        let keep = scan_bytes.len().min(SCAN_OVERLAP_BYTES);
293        self.scan_overlaps.insert(
294            cursor_key.to_string(),
295            scan_bytes[scan_bytes.len().saturating_sub(keep)..].to_vec(),
296        );
297        matches
298    }
299}
300
301fn find_match(
302    pattern: &WatchPattern,
303    text: &str,
304    min_end_exclusive: usize,
305) -> Option<(usize, usize, String)> {
306    match pattern {
307        WatchPattern::Substring(needle) => {
308            if needle.is_empty() {
309                return None;
310            }
311            let mut search_start = min_end_exclusive.saturating_sub(needle.len().saturating_sub(1));
312            while search_start > 0 && !text.is_char_boundary(search_start) {
313                search_start -= 1;
314            }
315            text.get(search_start..).and_then(|tail| {
316                tail.find(needle).and_then(|relative_start| {
317                    let start = search_start + relative_start;
318                    let end = start + needle.len();
319                    (end > min_end_exclusive).then(|| (start, end, needle.clone()))
320                })
321            })
322        }
323        WatchPattern::Regex(regex) => regex
324            .find_iter(text)
325            .find(|m| m.end() > min_end_exclusive)
326            .map(|m| (m.start(), m.end(), m.as_str().to_string())),
327    }
328}
329
330fn context_snippet(text: &str, start: usize, end: usize) -> String {
331    let before_start = text[..start]
332        .char_indices()
333        .rev()
334        .nth(CONTEXT_BEFORE)
335        .map(|(idx, _)| idx)
336        .unwrap_or(0);
337    let after_end = text[end..]
338        .char_indices()
339        .nth(CONTEXT_AFTER)
340        .map(|(idx, _)| end + idx)
341        .unwrap_or(text.len());
342    text[before_start..after_end].replace('\r', "")
343}
344
345#[cfg(test)]
346mod tests {
347    use super::*;
348
349    #[test]
350    fn once_watch_self_removes_after_match() {
351        let mut registry = WatchRegistry::default();
352        let task_id = "bash-1".to_string();
353        registry
354            .register(
355                task_id.clone(),
356                WatchPattern::Substring("READY".into()),
357                true,
358            )
359            .unwrap();
360        assert_eq!(registry.scan_new_bytes(&task_id, b"READY\n").len(), 1);
361        assert_eq!(registry.active_count(&task_id), 0);
362    }
363
364    #[test]
365    fn sticky_watch_fires_multiple_times() {
366        let mut registry = WatchRegistry::default();
367        let task_id = "bash-1".to_string();
368        registry
369            .register(
370                task_id.clone(),
371                WatchPattern::Substring("READY".into()),
372                false,
373            )
374            .unwrap();
375        assert_eq!(registry.scan_new_bytes(&task_id, b"READY\n").len(), 1);
376        assert_eq!(registry.scan_new_bytes(&task_id, b"READY\n").len(), 1);
377        assert_eq!(registry.active_count(&task_id), 1);
378    }
379
380    #[test]
381    fn cap_8_watches_per_task_rejects_9th() {
382        let mut registry = WatchRegistry::default();
383        for _ in 0..8 {
384            registry
385                .register("bash-1".into(), WatchPattern::Substring("x".into()), true)
386                .unwrap();
387        }
388        assert_eq!(
389            registry.register("bash-1".into(), WatchPattern::Substring("x".into()), true),
390            Err("too_many_watches")
391        );
392    }
393
394    #[test]
395    fn regex_pattern_matches_with_capture() {
396        let mut registry = WatchRegistry::default();
397        let task_id = "bash-1".to_string();
398        registry
399            .register(
400                task_id.clone(),
401                WatchPattern::regex("port (\\d+)").unwrap(),
402                true,
403            )
404            .unwrap();
405        let hits = registry.scan_new_bytes(&task_id, b"listening on port 3000\n");
406        assert_eq!(hits[0].match_text, "port 3000");
407    }
408
409    #[test]
410    fn substring_pattern_can_span_scans() {
411        let mut registry = WatchRegistry::default();
412        let task_id = "bash-1".to_string();
413        registry
414            .register(
415                task_id.clone(),
416                WatchPattern::Substring("READY".into()),
417                true,
418            )
419            .unwrap();
420
421        assert!(registry.scan_new_bytes(&task_id, b"RE").is_empty());
422        let hits = registry.scan_new_bytes(&task_id, b"ADY\n");
423
424        assert_eq!(hits.len(), 1);
425        assert_eq!(hits[0].match_text, "READY");
426        assert_eq!(hits[0].match_offset, 0);
427    }
428
429    #[test]
430    fn regex_pattern_can_span_scans() {
431        let mut registry = WatchRegistry::default();
432        let task_id = "bash-1".to_string();
433        registry
434            .register(
435                task_id.clone(),
436                WatchPattern::regex("ready: \\d{4}").unwrap(),
437                true,
438            )
439            .unwrap();
440
441        assert!(registry
442            .scan_new_bytes(&task_id, b"prefix ready: 4")
443            .is_empty());
444        let hits = registry.scan_new_bytes(&task_id, b"242\n");
445
446        assert_eq!(hits.len(), 1);
447        assert_eq!(hits[0].match_text, "ready: 4242");
448        assert_eq!(hits[0].match_offset, 7);
449    }
450
451    #[test]
452    fn overlap_does_not_repeat_fully_previous_match() {
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
463        assert_eq!(registry.scan_new_bytes(&task_id, b"READY").len(), 1);
464        assert!(registry.scan_new_bytes(&task_id, b"\n").is_empty());
465    }
466}
aft/bash_background/watches.rs

aft/bash_background/
watches.rs
