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

aft/bash_background/
watches.rs
