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
37#[derive(Debug, Clone, Serialize)]
38pub struct PatternMatch {
39 pub watch_id: String,
40 pub task_id: String,
41 pub match_text: String,
42 pub match_offset: u64,
43 pub context: String,
44 pub once: bool,
45}
46
47#[derive(Debug, Default)]
48pub struct WatchRegistry {
49 watches: HashMap<String, Vec<WatchSpec>>,
50 scan_cursors: HashMap<String, u64>,
51 scan_overlaps: HashMap<String, Vec<u8>>,
52 controlled_tasks: HashSet<String>,
53 matched_tasks: HashSet<String>,
54 next_watch: u64,
55}
56
57impl WatchRegistry {
58 pub fn register(
59 &mut self,
60 task_id: String,
61 pattern: WatchPattern,
62 once: bool,
63 ) -> Result<String, &'static str> {
64 let watches = self.watches.entry(task_id.clone()).or_default();
65 if watches.len() >= MAX_WATCHES_PER_TASK {
66 return Err("too_many_watches");
67 }
68 self.controlled_tasks.insert(task_id.clone());
69 self.next_watch = self.next_watch.wrapping_add(1);
70 let watch_id = format!("watch-{:08x}", self.next_watch);
71 watches.push(WatchSpec {
72 watch_id: watch_id.clone(),
73 task_id,
74 pattern,
75 once,
76 });
77 Ok(watch_id)
78 }
79
80 pub fn unregister(&mut self, task_id: &str, watch_id: &str) {
81 if let Some(watches) = self.watches.get_mut(task_id) {
82 watches.retain(|watch| watch.watch_id != watch_id);
83 if watches.is_empty() {
84 self.watches.remove(task_id);
85 }
86 }
87 }
88
89 pub fn clear_task(&mut self, task_id: &str) {
90 self.watches.remove(task_id);
91 self.controlled_tasks.remove(task_id);
92 self.matched_tasks.remove(task_id);
93 let prefix = format!("{task_id}:");
94 self.scan_cursors
95 .retain(|key, _| key != task_id && !key.starts_with(&prefix));
96 self.scan_overlaps
97 .retain(|key, _| key != task_id && !key.starts_with(&prefix));
98 }
99
100 pub fn has_controlled_task(&self, task_id: &str) -> bool {
101 self.controlled_tasks.contains(task_id)
102 }
103
104 pub fn has_matched_task(&self, task_id: &str) -> bool {
105 self.matched_tasks.contains(task_id)
106 }
107
108 pub fn active_count(&self, task_id: &str) -> usize {
109 self.watches.get(task_id).map_or(0, Vec::len)
110 }
111
112 pub fn prime_file_cursor(&mut self, cursor_key: &str, file: &ValidatedArtifact) {
113 if self.scan_cursors.contains_key(cursor_key) {
114 return;
115 }
116 let len = file.len().unwrap_or(0);
117 self.scan_cursors.insert(cursor_key.to_string(), len);
118 }
119
120 pub fn set_file_cursor(&mut self, cursor_key: &str, offset: u64) {
121 self.scan_cursors.insert(cursor_key.to_string(), offset);
122 self.scan_overlaps.remove(cursor_key);
123 }
124
125 pub fn scan_file_new_bytes(
126 &mut self,
127 cursor_key: &str,
128 task_id: &str,
129 file: &mut ValidatedArtifact,
130 ) -> Vec<PatternMatch> {
131 if self.active_count(task_id) == 0 {
132 return Vec::new();
133 }
134 let cursor = self
135 .scan_cursors
136 .get(cursor_key)
137 .copied()
138 .unwrap_or_else(|| {
139 file.len().unwrap_or(0)
141 });
142 if file.seek(SeekFrom::Start(cursor)).is_err() {
143 return Vec::new();
144 }
145 let mut bytes = Vec::new();
146 if file.read_to_end(&mut bytes).is_err() || bytes.is_empty() {
147 self.scan_cursors.insert(cursor_key.to_string(), cursor);
148 return Vec::new();
149 }
150 let next = cursor.saturating_add(bytes.len() as u64);
151 self.scan_cursors.insert(cursor_key.to_string(), next);
152 self.scan_new_bytes_at(cursor_key, task_id, &bytes, cursor)
153 }
154
155 pub fn scan_new_bytes(&mut self, task_id: &str, bytes: &[u8]) -> Vec<PatternMatch> {
156 let base = self.scan_cursors.get(task_id).copied().unwrap_or(0);
157 self.scan_cursors
158 .insert(task_id.to_string(), base.saturating_add(bytes.len() as u64));
159 self.scan_new_bytes_at(task_id, task_id, bytes, base)
160 }
161
162 fn scan_new_bytes_at(
163 &mut self,
164 cursor_key: &str,
165 task_id: &str,
166 bytes: &[u8],
167 base_offset: u64,
168 ) -> Vec<PatternMatch> {
169 let Some(watches) = self.watches.get(task_id).cloned() else {
170 return Vec::new();
171 };
172 let overlap = self
173 .scan_overlaps
174 .get(cursor_key)
175 .cloned()
176 .unwrap_or_default();
177 let prefix_len = overlap.len();
178 let mut scan_bytes = Vec::with_capacity(prefix_len.saturating_add(bytes.len()));
179 scan_bytes.extend_from_slice(&overlap);
180 scan_bytes.extend_from_slice(bytes);
181 let text = String::from_utf8_lossy(&scan_bytes);
182 let scan_base_offset = base_offset.saturating_sub(prefix_len as u64);
183 let mut matches = Vec::new();
184 let mut remove_once = Vec::new();
185 for watch in watches {
186 if let Some((start, end, matched)) = find_match(&watch.pattern, &text, prefix_len) {
187 self.matched_tasks.insert(task_id.to_string());
188 matches.push(PatternMatch {
189 watch_id: watch.watch_id.clone(),
190 task_id: watch.task_id.clone(),
191 match_text: matched,
192 match_offset: scan_base_offset.saturating_add(start as u64),
193 context: context_snippet(&text, start, end),
194 once: watch.once,
195 });
196 if watch.once {
197 remove_once.push(watch.watch_id);
198 }
199 }
200 }
201 for watch_id in remove_once {
202 self.unregister(task_id, &watch_id);
203 }
204 let keep = scan_bytes.len().min(SCAN_OVERLAP_BYTES);
205 self.scan_overlaps.insert(
206 cursor_key.to_string(),
207 scan_bytes[scan_bytes.len().saturating_sub(keep)..].to_vec(),
208 );
209 matches
210 }
211}
212
213fn find_match(
214 pattern: &WatchPattern,
215 text: &str,
216 min_end_exclusive: usize,
217) -> Option<(usize, usize, String)> {
218 match pattern {
219 WatchPattern::Substring(needle) => {
220 if needle.is_empty() {
221 return None;
222 }
223 let mut search_start = min_end_exclusive.saturating_sub(needle.len().saturating_sub(1));
224 while search_start > 0 && !text.is_char_boundary(search_start) {
225 search_start -= 1;
226 }
227 text.get(search_start..).and_then(|tail| {
228 tail.find(needle).and_then(|relative_start| {
229 let start = search_start + relative_start;
230 let end = start + needle.len();
231 (end > min_end_exclusive).then(|| (start, end, needle.clone()))
232 })
233 })
234 }
235 WatchPattern::Regex(regex) => regex
236 .find_iter(text)
237 .find(|m| m.end() > min_end_exclusive)
238 .map(|m| (m.start(), m.end(), m.as_str().to_string())),
239 }
240}
241
242fn context_snippet(text: &str, start: usize, end: usize) -> String {
243 let before_start = text[..start]
244 .char_indices()
245 .rev()
246 .nth(CONTEXT_BEFORE)
247 .map(|(idx, _)| idx)
248 .unwrap_or(0);
249 let after_end = text[end..]
250 .char_indices()
251 .nth(CONTEXT_AFTER)
252 .map(|(idx, _)| end + idx)
253 .unwrap_or(text.len());
254 text[before_start..after_end].replace('\r', "")
255}
256
257#[cfg(test)]
258mod tests {
259 use super::*;
260
261 #[test]
262 fn once_watch_self_removes_after_match() {
263 let mut registry = WatchRegistry::default();
264 let task_id = "bash-1".to_string();
265 registry
266 .register(
267 task_id.clone(),
268 WatchPattern::Substring("READY".into()),
269 true,
270 )
271 .unwrap();
272 assert_eq!(registry.scan_new_bytes(&task_id, b"READY\n").len(), 1);
273 assert_eq!(registry.active_count(&task_id), 0);
274 }
275
276 #[test]
277 fn sticky_watch_fires_multiple_times() {
278 let mut registry = WatchRegistry::default();
279 let task_id = "bash-1".to_string();
280 registry
281 .register(
282 task_id.clone(),
283 WatchPattern::Substring("READY".into()),
284 false,
285 )
286 .unwrap();
287 assert_eq!(registry.scan_new_bytes(&task_id, b"READY\n").len(), 1);
288 assert_eq!(registry.scan_new_bytes(&task_id, b"READY\n").len(), 1);
289 assert_eq!(registry.active_count(&task_id), 1);
290 }
291
292 #[test]
293 fn cap_8_watches_per_task_rejects_9th() {
294 let mut registry = WatchRegistry::default();
295 for _ in 0..8 {
296 registry
297 .register("bash-1".into(), WatchPattern::Substring("x".into()), true)
298 .unwrap();
299 }
300 assert_eq!(
301 registry.register("bash-1".into(), WatchPattern::Substring("x".into()), true),
302 Err("too_many_watches")
303 );
304 }
305
306 #[test]
307 fn regex_pattern_matches_with_capture() {
308 let mut registry = WatchRegistry::default();
309 let task_id = "bash-1".to_string();
310 registry
311 .register(
312 task_id.clone(),
313 WatchPattern::regex("port (\\d+)").unwrap(),
314 true,
315 )
316 .unwrap();
317 let hits = registry.scan_new_bytes(&task_id, b"listening on port 3000\n");
318 assert_eq!(hits[0].match_text, "port 3000");
319 }
320
321 #[test]
322 fn substring_pattern_can_span_scans() {
323 let mut registry = WatchRegistry::default();
324 let task_id = "bash-1".to_string();
325 registry
326 .register(
327 task_id.clone(),
328 WatchPattern::Substring("READY".into()),
329 true,
330 )
331 .unwrap();
332
333 assert!(registry.scan_new_bytes(&task_id, b"RE").is_empty());
334 let hits = registry.scan_new_bytes(&task_id, b"ADY\n");
335
336 assert_eq!(hits.len(), 1);
337 assert_eq!(hits[0].match_text, "READY");
338 assert_eq!(hits[0].match_offset, 0);
339 }
340
341 #[test]
342 fn regex_pattern_can_span_scans() {
343 let mut registry = WatchRegistry::default();
344 let task_id = "bash-1".to_string();
345 registry
346 .register(
347 task_id.clone(),
348 WatchPattern::regex("ready: \\d{4}").unwrap(),
349 true,
350 )
351 .unwrap();
352
353 assert!(registry
354 .scan_new_bytes(&task_id, b"prefix ready: 4")
355 .is_empty());
356 let hits = registry.scan_new_bytes(&task_id, b"242\n");
357
358 assert_eq!(hits.len(), 1);
359 assert_eq!(hits[0].match_text, "ready: 4242");
360 assert_eq!(hits[0].match_offset, 7);
361 }
362
363 #[test]
364 fn overlap_does_not_repeat_fully_previous_match() {
365 let mut registry = WatchRegistry::default();
366 let task_id = "bash-1".to_string();
367 registry
368 .register(
369 task_id.clone(),
370 WatchPattern::Substring("READY".into()),
371 false,
372 )
373 .unwrap();
374
375 assert_eq!(registry.scan_new_bytes(&task_id, b"READY").len(), 1);
376 assert!(registry.scan_new_bytes(&task_id, b"\n").is_empty());
377 }
378}