atman_runtime/
watch.rs

1use std::collections::HashMap;
2use std::sync::{Arc, Mutex};
3use std::time::Duration;
4
5use tokio::sync::Notify;
6use tokio_util::sync::CancellationToken;
7
8use crate::error::RuntimeError;
9use crate::message::Message;
10use crate::tool::{BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
11use crate::value::Value;
12
13pub type WatcherId = String;
14
15#[derive(Debug, Clone)]
16pub enum WatchSource {
17    Terminal { handle: String },
18    Bash { handle: String },
19    Agent { handle: String },
20}
21
22impl WatchSource {
23    pub fn kind_str(&self) -> &'static str {
24        match self {
25            Self::Terminal { .. } => "terminal",
26            Self::Bash { .. } => "bash",
27            Self::Agent { .. } => "agent",
28        }
29    }
30
31    pub fn handle(&self) -> &str {
32        match self {
33            Self::Terminal { handle } | Self::Bash { handle } | Self::Agent { handle } => handle,
34        }
35    }
36}
37
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub enum WatchMode {
40    Once,
41    Persist,
42}
43
44#[derive(Debug, Clone)]
45pub struct WatchEvent {
46    pub watcher_id: WatcherId,
47    pub source: WatchSource,
48    pub pattern: String,
49    pub row: Option<u16>,
50    pub col: Option<u16>,
51    pub text: String,
52    pub timestamp: chrono::DateTime<chrono::Utc>,
53    pub timed_out: bool,
54    pub exited: bool,
55    pub timeout: Duration,
56}
57
58#[derive(Debug, Clone)]
59pub enum WatchResult {
60    Matched {
61        row: Option<u16>,
62        col: Option<u16>,
63        text: String,
64    },
65    SourceExited,
66    Cancelled,
67}
68
69struct Watcher {
70    id: WatcherId,
71    source: WatchSource,
72    pattern: String,
73    mode: WatchMode,
74    timeout: Duration,
75    created_at: chrono::DateTime<chrono::Utc>,
76    cancel: CancellationToken,
77}
78
79#[derive(Debug, Clone)]
80pub struct WatcherInfo {
81    pub id: WatcherId,
82    pub source: WatchSource,
83    pub pattern: String,
84    pub mode: WatchMode,
85    pub age: Duration,
86    pub timeout: Duration,
87}
88
89pub struct WatchHub {
90    watchers: Mutex<HashMap<WatcherId, Watcher>>,
91    pending_events: Mutex<Vec<WatchEvent>>,
92    notify: Notify,
93}
94
95impl Default for WatchHub {
96    fn default() -> Self {
97        Self::new()
98    }
99}
100
101impl WatchHub {
102    pub fn new() -> Self {
103        Self {
104            watchers: Mutex::new(HashMap::new()),
105            pending_events: Mutex::new(Vec::new()),
106            notify: Notify::new(),
107        }
108    }
109
110    pub fn register(
111        &self,
112        source: WatchSource,
113        pattern: String,
114        mode: WatchMode,
115        timeout: Duration,
116    ) -> WatcherId {
117        let id = format!("w_{}", uuid::Uuid::now_v7().simple());
118        let cancel = CancellationToken::new();
119        let watcher = Watcher {
120            id: id.clone(),
121            source,
122            pattern,
123            mode,
124            timeout,
125            created_at: chrono::Utc::now(),
126            cancel,
127        };
128        self.watchers.lock().unwrap().insert(id.clone(), watcher);
129        id
130    }
131
132    pub fn unregister(&self, id: &WatcherId) -> bool {
133        if let Some(w) = self.watchers.lock().unwrap().remove(id) {
134            w.cancel.cancel();
135            true
136        } else {
137            false
138        }
139    }
140
141    pub fn has_active_watchers(&self) -> bool {
142        !self.watchers.lock().unwrap().is_empty()
143    }
144
145    pub fn list_active(&self) -> Vec<WatcherInfo> {
146        let now = chrono::Utc::now();
147        self.watchers
148            .lock()
149            .unwrap()
150            .values()
151            .map(|w| WatcherInfo {
152                id: w.id.clone(),
153                source: w.source.clone(),
154                pattern: w.pattern.clone(),
155                mode: w.mode,
156                age: (now - w.created_at).to_std().unwrap_or(Duration::ZERO),
157                timeout: w.timeout,
158            })
159            .collect()
160    }
161
162    pub fn list_watchers_for_handle(&self, handle: &str) -> Vec<WatcherInfo> {
163        self.list_active()
164            .into_iter()
165            .filter(|w| w.source.handle() == handle)
166            .collect()
167    }
168
169    pub fn get_cancel(&self, id: &WatcherId) -> CancellationToken {
170        self.watchers
171            .lock()
172            .unwrap()
173            .get(id)
174            .map(|w| w.cancel.clone())
175            .unwrap_or_default()
176    }
177
178    pub fn enqueue_event(&self, event: WatchEvent) {
179        let should_remove = {
180            let watchers = self.watchers.lock().unwrap();
181            watchers
182                .get(&event.watcher_id)
183                .map(|w| matches!(w.mode, WatchMode::Once))
184                .unwrap_or(false)
185        };
186        if should_remove {
187            if let Some(w) = self.watchers.lock().unwrap().remove(&event.watcher_id) {
188                w.cancel.cancel();
189            }
190        }
191        self.pending_events.lock().unwrap().push(event);
192        self.notify.notify_one();
193    }
194
195    pub async fn wait_for_event(&self, timeout: Duration) -> Option<WatchEvent> {
196        if let Some(evt) = self.pending_events.lock().unwrap().pop() {
197            return Some(evt);
198        }
199        if !self.has_active_watchers() {
200            return None;
201        }
202        let _ = tokio::time::timeout(timeout, self.notify.notified()).await;
203        self.pending_events.lock().unwrap().pop()
204    }
205}
206
207pub trait Watchable: Send + Sync {
208    fn watch_output(
209        self: Arc<Self>,
210        pattern: String,
211        cancel: CancellationToken,
212    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = WatchResult> + Send>>;
213}
214
215pub async fn handle_watch_result(
216    hub: Arc<WatchHub>,
217    wid: WatcherId,
218    source: WatchSource,
219    pattern: String,
220    timeout: Duration,
221    result: Result<WatchResult, tokio::time::error::Elapsed>,
222) {
223    match result {
224        Ok(WatchResult::Matched { row, col, text }) => {
225            hub.enqueue_event(WatchEvent {
226                watcher_id: wid,
227                source,
228                pattern,
229                row,
230                col,
231                text,
232                timestamp: chrono::Utc::now(),
233                timed_out: false,
234                exited: false,
235                timeout,
236            });
237        }
238        Ok(WatchResult::SourceExited) => {
239            hub.enqueue_event(WatchEvent {
240                watcher_id: wid.clone(),
241                source,
242                pattern,
243                row: None,
244                col: None,
245                text: String::new(),
246                timestamp: chrono::Utc::now(),
247                timed_out: false,
248                exited: true,
249                timeout,
250            });
251            hub.unregister(&wid);
252        }
253        Ok(WatchResult::Cancelled) => {
254            hub.unregister(&wid);
255        }
256        Err(_) => {
257            hub.enqueue_event(WatchEvent {
258                watcher_id: wid,
259                source,
260                pattern,
261                row: None,
262                col: None,
263                text: String::new(),
264                timestamp: chrono::Utc::now(),
265                timed_out: true,
266                exited: false,
267                timeout,
268            });
269        }
270    }
271}
272
273pub fn format_watch_event_text(evt: &WatchEvent) -> String {
274    let kind = evt.source.kind_str();
275    let handle = evt.source.handle();
276    if evt.exited {
277        format!(
278            "[watcher {}] {} '{}' has already exited. Pattern '{}' will never match.",
279            evt.watcher_id, kind, handle, evt.pattern
280        )
281    } else if evt.timed_out {
282        format!(
283            "[watcher {}] {} '{}' pattern '{}' not detected in {}s. \
284             Consider using {}.capture or {}.output to check current state.",
285            evt.watcher_id,
286            kind,
287            handle,
288            evt.pattern,
289            evt.timeout.as_secs(),
290            kind,
291            kind
292        )
293    } else {
294        format!(
295            "[watcher {}] {} '{}' matched '{}' at row {:?}, col {:?}: {}",
296            evt.watcher_id, kind, handle, evt.pattern, evt.row, evt.col, evt.text
297        )
298    }
299}
300
301fn watch_source_to_value(src: &WatchSource) -> Value {
302    Value::Struct(vec![
303        ("kind".into(), Value::Str(src.kind_str().into())),
304        ("handle".into(), Value::Str(src.handle().into())),
305    ])
306}
307
308fn extract_string(args: &ToolArgs, name: &str, pos: usize) -> Result<String, RuntimeError> {
309    let value = match args.named(name) {
310        Some(v) => v,
311        None => args.positional(pos)?,
312    };
313    match value {
314        Value::Str(s) => Ok(s.clone()),
315        other => Err(RuntimeError::TypeMismatch {
316            expected: "string".into(),
317            actual: other.kind_name().into(),
318        }),
319    }
320}
321
322fn extract_optional_string(args: &ToolArgs, name: &str) -> Option<String> {
323    args.named(name).and_then(|v| {
324        if let Value::Str(s) = v {
325            Some(s.clone())
326        } else {
327            None
328        }
329    })
330}
331
332fn extract_optional_int(args: &ToolArgs, name: &str) -> Option<i64> {
333    args.named(name).and_then(|v| {
334        if let Value::Int(i) = v {
335            Some(*i)
336        } else {
337            None
338        }
339    })
340}
341
342pub struct Watch;
343impl Tool for Watch {
344    fn name(&self) -> &str {
345        "watch"
346    }
347    fn tier(&self) -> Tier {
348        Tier::Four
349    }
350    fn description(&self) -> Option<&str> {
351        Some(
352            "Register a background watcher on any running task (terminal, bash, or agent).\n\
353             When pattern appears in the task's output, the agent is woken up.\n\n\
354             Non-blocking: returns watcher_id immediately.\n\
355             mode: \"once\" (default) or \"persist\".\n\
356             timeout_ms: default 120000 (120s), max 600000 (10min).",
357        )
358    }
359    fn input_schema(&self) -> serde_json::Value {
360        serde_json::json!({
361            "type": "object",
362            "properties": {
363                "handle": {"type": "string"},
364                "pattern": {"type": "string"},
365                "mode": {"type": "string", "enum": ["once", "persist"], "default": "once"},
366                "timeout_ms": {"type": "integer", "default": 120000}
367            },
368            "required": ["handle", "pattern"]
369        })
370    }
371    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
372        Box::pin(async move {
373            let handle = extract_string(&args, "handle", 0)?;
374            let pattern = extract_string(&args, "pattern", 1)?;
375            let mode = match extract_optional_string(&args, "mode").as_deref() {
376                Some("persist") => WatchMode::Persist,
377                _ => WatchMode::Once,
378            };
379            let timeout_ms = extract_optional_int(&args, "timeout_ms")
380                .unwrap_or(120_000)
381                .clamp(1_000, 600_000) as u64;
382            let session_id = ctx.session_id.clone().unwrap_or_else(|| "anon".into());
383
384            let task_registry = ctx.task_registry.clone().ok_or_else(|| {
385                RuntimeError::ToolFailed("watch: task registry not available".into())
386            })?;
387            let snapshot = task_registry.lookup_by_handle(&handle).ok_or_else(|| {
388                RuntimeError::ToolFailed(format!("watch: handle '{handle}' not found"))
389            })?;
390
391            let watch_hub = ctx
392                .watch_hub
393                .clone()
394                .ok_or_else(|| RuntimeError::ToolFailed("watch: watch hub not available".into()))?;
395
396            let source = match snapshot.kind {
397                crate::task_registry::TaskKind::Terminal => WatchSource::Terminal {
398                    handle: handle.clone(),
399                },
400                crate::task_registry::TaskKind::Bash => WatchSource::Bash {
401                    handle: handle.clone(),
402                },
403                crate::task_registry::TaskKind::Flow => WatchSource::Agent {
404                    handle: handle.clone(),
405                },
406            };
407
408            let watcher_id = watch_hub.register(
409                source.clone(),
410                pattern.clone(),
411                mode,
412                Duration::from_millis(timeout_ms),
413            );
414
415            match snapshot.kind {
416                crate::task_registry::TaskKind::Terminal => {
417                    let reg = ctx.term_registry.clone().ok_or_else(|| {
418                        RuntimeError::ToolFailed("watch: terminal registry not available".into())
419                    })?;
420                    let entry = reg.lookup(&handle, &session_id)?;
421                    let watchable: Arc<dyn Watchable> = Arc::clone(&entry) as Arc<dyn Watchable>;
422                    spawn_watcher(
423                        watch_hub,
424                        watcher_id.clone(),
425                        source,
426                        pattern,
427                        mode,
428                        Duration::from_millis(timeout_ms),
429                        watchable,
430                    );
431                }
432                crate::task_registry::TaskKind::Bash => {
433                    let reg = ctx.bg_registry.clone().ok_or_else(|| {
434                        RuntimeError::ToolFailed("watch: bash registry not available".into())
435                    })?;
436                    let entry = reg.lookup(&handle, &session_id)?;
437                    let watchable: Arc<dyn Watchable> = Arc::clone(&entry) as Arc<dyn Watchable>;
438                    spawn_watcher(
439                        watch_hub,
440                        watcher_id.clone(),
441                        source,
442                        pattern,
443                        mode,
444                        Duration::from_millis(timeout_ms),
445                        watchable,
446                    );
447                }
448                crate::task_registry::TaskKind::Flow => {
449                    let reg = ctx.flow_registry.clone().ok_or_else(|| {
450                        RuntimeError::ToolFailed("watch: agent registry not available".into())
451                    })?;
452                    let entry = reg.lookup(&handle)?;
453                    let watchable: Arc<dyn Watchable> = Arc::clone(&entry) as Arc<dyn Watchable>;
454                    spawn_watcher(
455                        watch_hub,
456                        watcher_id.clone(),
457                        source,
458                        pattern,
459                        mode,
460                        Duration::from_millis(timeout_ms),
461                        watchable,
462                    );
463                }
464            }
465
466            Ok(Value::Struct(vec![(
467                "watcher_id".into(),
468                Value::Str(watcher_id),
469            )]))
470        })
471    }
472}
473
474fn spawn_watcher(
475    hub: Arc<WatchHub>,
476    wid: WatcherId,
477    source: WatchSource,
478    pattern: String,
479    mode: WatchMode,
480    timeout: Duration,
481    watchable: Arc<dyn Watchable>,
482) {
483    let cancel = hub.get_cancel(&wid);
484    let w = Arc::clone(&watchable);
485    let pat = pattern.clone();
486    tokio::spawn(async move {
487        loop {
488            let result = tokio::time::timeout(
489                timeout,
490                Arc::clone(&w).watch_output(pat.clone(), cancel.clone()),
491            )
492            .await;
493            let matched = matches!(result, Ok(WatchResult::Matched { .. }));
494            handle_watch_result(
495                Arc::clone(&hub),
496                wid.clone(),
497                source.clone(),
498                pattern.clone(),
499                timeout,
500                result,
501            )
502            .await;
503            if !matched || mode == WatchMode::Once {
504                break;
505            }
506        }
507        if mode == WatchMode::Persist {
508            hub.unregister(&wid);
509        }
510    });
511}
512
513pub struct WatcherList;
514impl Tool for WatcherList {
515    fn name(&self) -> &str {
516        "watcher.list"
517    }
518    fn tier(&self) -> Tier {
519        Tier::Zero
520    }
521    fn description(&self) -> Option<&str> {
522        Some("List all active background watchers with their source, pattern, mode, and age.")
523    }
524    fn input_schema(&self) -> serde_json::Value {
525        serde_json::json!({"type": "object", "properties": {}})
526    }
527    fn call<'a>(&'a self, _args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
528        Box::pin(async move {
529            let hub = ctx.watch_hub.clone().ok_or_else(|| {
530                RuntimeError::ToolFailed("watcher.list: watch hub not available".into())
531            })?;
532            let list = hub.list_active();
533            let items: Vec<Value> = list
534                .iter()
535                .map(|w| {
536                    Value::Struct(vec![
537                        ("watcher_id".into(), Value::Str(w.id.clone())),
538                        ("source".into(), watch_source_to_value(&w.source)),
539                        ("pattern".into(), Value::Str(w.pattern.clone())),
540                        (
541                            "mode".into(),
542                            Value::Str(
543                                match w.mode {
544                                    WatchMode::Once => "once",
545                                    WatchMode::Persist => "persist",
546                                }
547                                .into(),
548                            ),
549                        ),
550                        ("age_ms".into(), Value::Int(w.age.as_millis() as i64)),
551                        (
552                            "timeout_ms".into(),
553                            Value::Int(w.timeout.as_millis() as i64),
554                        ),
555                    ])
556                })
557                .collect();
558            Ok(Value::List(items))
559        })
560    }
561}
562
563pub struct WatcherUnwatch;
564impl Tool for WatcherUnwatch {
565    fn name(&self) -> &str {
566        "watcher.unwatch"
567    }
568    fn tier(&self) -> Tier {
569        Tier::Zero
570    }
571    fn description(&self) -> Option<&str> {
572        Some("Cancel a background watcher by id. Works for any source (terminal/bash/agent).")
573    }
574    fn input_schema(&self) -> serde_json::Value {
575        serde_json::json!({
576            "type": "object",
577            "properties": {
578                "watcher_id": {"type": "string"}
579            },
580            "required": ["watcher_id"]
581        })
582    }
583    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
584        Box::pin(async move {
585            let id = extract_string(&args, "watcher_id", 0)?;
586            let hub = ctx.watch_hub.clone().ok_or_else(|| {
587                RuntimeError::ToolFailed("watcher.unwatch: watch hub not available".into())
588            })?;
589            if hub.unregister(&id) {
590                Ok(Value::Unit)
591            } else {
592                Err(RuntimeError::ToolFailed(format!(
593                    "watcher.unwatch: '{id}' not found"
594                )))
595            }
596        })
597    }
598}
599
600pub struct WaitForWatcher;
601impl Tool for WaitForWatcher {
602    fn name(&self) -> &str {
603        "wait_for_watcher"
604    }
605    fn tier(&self) -> Tier {
606        Tier::Zero
607    }
608    fn description(&self) -> Option<&str> {
609        Some(
610            "Block until a registered watcher fires, or timeout.\n\
611             If no active watchers exist, returns immediately (Unit).\n\
612             Returns Str with event details on match, or Unit on timeout/no watchers.",
613        )
614    }
615    fn input_schema(&self) -> serde_json::Value {
616        serde_json::json!({
617            "type": "object",
618            "properties": {
619                "timeout_ms": {"type": "integer", "default": 30000}
620            }
621        })
622    }
623    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
624        Box::pin(async move {
625            let timeout_ms = extract_optional_int(&args, "timeout_ms")
626                .unwrap_or(30_000)
627                .clamp(100, 300_000) as u64;
628            let hub = ctx.watch_hub.clone().ok_or_else(|| {
629                RuntimeError::ToolFailed("wait_for_watcher: watch hub not available".into())
630            })?;
631            match hub.wait_for_event(Duration::from_millis(timeout_ms)).await {
632                Some(evt) => Ok(Value::Message(Message::system_text(
633                    ctx.turn_id
634                        .clone()
635                        .unwrap_or_else(crate::event::TurnId::now),
636                    format_watch_event_text(&evt),
637                ))),
638                None => Ok(Value::Unit),
639            }
640        })
641    }
642}
643
644pub struct HasPendingInjections;
645impl Tool for HasPendingInjections {
646    fn name(&self) -> &str {
647        "has_pending_injections"
648    }
649    fn tier(&self) -> Tier {
650        Tier::Zero
651    }
652    fn description(&self) -> Option<&str> {
653        Some(
654            "Check if there are pending user injections for the current turn.\n\
655             Returns true if any exist — the agent loop should continue so they get drained.",
656        )
657    }
658    fn input_schema(&self) -> serde_json::Value {
659        serde_json::json!({"type": "object", "properties": {}})
660    }
661    fn call<'a>(&'a self, _args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
662        Box::pin(async move {
663            if ctx
664                .agent_entry
665                .as_ref()
666                .is_some_and(|entry| !entry.pending_injections.lock().unwrap().is_empty())
667            {
668                return Ok(Value::Bool(true));
669            }
670            let Some(session) = ctx.session_runtime.as_ref() else {
671                return if ctx.agent_entry.is_some() {
672                    Ok(Value::Bool(false))
673                } else {
674                    Err(RuntimeError::ToolFailed(
675                        "has_pending_injections: no current flow".into(),
676                    ))
677                };
678            };
679            let turn_id = ctx.turn_id.clone().ok_or_else(|| {
680                RuntimeError::ToolFailed("has_pending_injections: no turn id".into())
681            })?;
682            let has_pending = session.list_pending_injections().iter().any(|inj| {
683                inj.turn_id == turn_id && inj.state == crate::injection::InjectionState::Pending
684            });
685            Ok(Value::Bool(has_pending))
686        })
687    }
688}
689
690#[cfg(test)]
691mod tests {
692    use super::*;
693
694    #[tokio::test]
695    async fn pending_injection_check_uses_the_current_flow_store() {
696        let session = Arc::new(crate::session::Session::open_ephemeral());
697        let turn_id = crate::event::TurnId::now();
698        session.begin_turn(crate::message::Message::user_text(turn_id.clone(), "root"));
699        session.enqueue_injection("root correction").unwrap();
700        let root_ctx = ToolCtx::new()
701            .with_anchors(Some(turn_id), None, None)
702            .with_session_runtime(session);
703        assert!(matches!(
704            HasPendingInjections
705                .call(ToolArgs::default(), &root_ctx)
706                .await
707                .unwrap(),
708            Value::Bool(true)
709        ));
710
711        let registry = crate::tools::agent_ctrl::FlowRegistry::new();
712        let entry = registry.create_entry(
713            "child".into(),
714            "goal".into(),
715            String::new(),
716            crate::event::FlowRunId::now(),
717        );
718        let child_ctx = ToolCtx::new().with_agent_entry(Arc::clone(&entry));
719        assert!(matches!(
720            HasPendingInjections
721                .call(ToolArgs::default(), &child_ctx)
722                .await
723                .unwrap(),
724            Value::Bool(false)
725        ));
726        entry
727            .pending_injections
728            .lock()
729            .unwrap()
730            .push(crate::injection::Injection::with_level(
731                crate::event::TurnId::now(),
732                "child correction",
733                crate::injection::InjectionLevel::L1Nudge,
734                None,
735            ));
736        assert!(matches!(
737            HasPendingInjections
738                .call(ToolArgs::default(), &child_ctx)
739                .await
740                .unwrap(),
741            Value::Bool(true)
742        ));
743    }
744
745    #[test]
746    fn watch_hub_register_and_unregister() {
747        let hub = WatchHub::new();
748        let id = hub.register(
749            WatchSource::Terminal {
750                handle: "t1".into(),
751            },
752            "$ ".into(),
753            WatchMode::Once,
754            Duration::from_secs(10),
755        );
756        assert!(hub.has_active_watchers());
757        assert!(hub.unregister(&id));
758        assert!(!hub.has_active_watchers());
759    }
760
761    #[test]
762    fn watch_hub_list_watchers_for_handle() {
763        let hub = WatchHub::new();
764        hub.register(
765            WatchSource::Terminal {
766                handle: "t1".into(),
767            },
768            "$ ".into(),
769            WatchMode::Once,
770            Duration::from_secs(10),
771        );
772        hub.register(
773            WatchSource::Terminal {
774                handle: "t1".into(),
775            },
776            "error".into(),
777            WatchMode::Persist,
778            Duration::from_secs(30),
779        );
780        hub.register(
781            WatchSource::Bash {
782                handle: "b1".into(),
783            },
784            "done".into(),
785            WatchMode::Once,
786            Duration::from_secs(10),
787        );
788        let t1_watchers = hub.list_watchers_for_handle("t1");
789        assert_eq!(t1_watchers.len(), 2);
790        let b1_watchers = hub.list_watchers_for_handle("b1");
791        assert_eq!(b1_watchers.len(), 1);
792        assert!(hub.list_watchers_for_handle("nonexistent").is_empty());
793    }
794
795    #[test]
796    fn watch_hub_enqueue_event_removes_once_mode() {
797        let hub = WatchHub::new();
798        let id = hub.register(
799            WatchSource::Bash {
800                handle: "b1".into(),
801            },
802            "done".into(),
803            WatchMode::Once,
804            Duration::from_secs(10),
805        );
806        hub.enqueue_event(WatchEvent {
807            watcher_id: id.clone(),
808            source: WatchSource::Bash {
809                handle: "b1".into(),
810            },
811            pattern: "done".into(),
812            row: Some(0),
813            col: Some(0),
814            text: "done".into(),
815            timestamp: chrono::Utc::now(),
816            timed_out: false,
817            exited: false,
818            timeout: Duration::from_secs(10),
819        });
820        assert!(
821            !hub.has_active_watchers(),
822            "once-mode watcher should be auto-removed after event"
823        );
824    }
825
826    #[test]
827    fn watch_hub_enqueue_event_keeps_persist_mode() {
828        let hub = WatchHub::new();
829        let id = hub.register(
830            WatchSource::Bash {
831                handle: "b1".into(),
832            },
833            "done".into(),
834            WatchMode::Persist,
835            Duration::from_secs(10),
836        );
837        hub.enqueue_event(WatchEvent {
838            watcher_id: id,
839            source: WatchSource::Bash {
840                handle: "b1".into(),
841            },
842            pattern: "done".into(),
843            row: Some(0),
844            col: Some(0),
845            text: "done".into(),
846            timestamp: chrono::Utc::now(),
847            timed_out: false,
848            exited: false,
849            timeout: Duration::from_secs(10),
850        });
851        assert!(
852            hub.has_active_watchers(),
853            "persist-mode watcher should remain after event"
854        );
855    }
856
857    #[test]
858    fn format_event_text_includes_watcher_id_and_pattern() {
859        let evt = WatchEvent {
860            watcher_id: "w_abc123".into(),
861            source: WatchSource::Terminal {
862                handle: "term_x".into(),
863            },
864            pattern: "$ ".into(),
865            row: Some(10),
866            col: Some(0),
867            text: "$ ls".into(),
868            timestamp: chrono::Utc::now(),
869            timed_out: false,
870            exited: false,
871            timeout: Duration::from_secs(120),
872        };
873        let text = format_watch_event_text(&evt);
874        assert!(text.contains("w_abc123"));
875        assert!(text.contains("terminal"));
876        assert!(text.contains("term_x"));
877        assert!(text.contains("$ "));
878    }
879
880    #[test]
881    fn format_event_text_timeout_includes_suggestion() {
882        let evt = WatchEvent {
883            watcher_id: "w_abc".into(),
884            source: WatchSource::Bash {
885                handle: "b1".into(),
886            },
887            pattern: "done".into(),
888            row: None,
889            col: None,
890            text: String::new(),
891            timestamp: chrono::Utc::now(),
892            timed_out: true,
893            exited: false,
894            timeout: Duration::from_secs(120),
895        };
896        let text = format_watch_event_text(&evt);
897        assert!(text.contains("not detected"));
898        assert!(text.contains("capture or"));
899    }
900
901    #[tokio::test]
902    async fn wait_for_event_returns_none_when_no_watchers() {
903        let hub = WatchHub::new();
904        let result = hub.wait_for_event(Duration::from_millis(100)).await;
905        assert!(result.is_none());
906    }
907
908    #[tokio::test]
909    async fn wait_for_event_returns_pending_event() {
910        let hub = Arc::new(WatchHub::new());
911        hub.register(
912            WatchSource::Terminal {
913                handle: "t1".into(),
914            },
915            "$ ".into(),
916            WatchMode::Once,
917            Duration::from_secs(10),
918        );
919        hub.enqueue_event(WatchEvent {
920            watcher_id: "w_test".into(),
921            source: WatchSource::Terminal {
922                handle: "t1".into(),
923            },
924            pattern: "$ ".into(),
925            row: Some(0),
926            col: Some(0),
927            text: "$ ".into(),
928            timestamp: chrono::Utc::now(),
929            timed_out: false,
930            exited: false,
931            timeout: Duration::from_secs(10),
932        });
933        let evt = hub.wait_for_event(Duration::from_millis(100)).await;
934        assert!(evt.is_some());
935        assert_eq!(evt.unwrap().pattern, "$ ");
936    }
937
938    #[tokio::test]
939    async fn wait_for_event_returns_source_exited_event() {
940        // When a watcher's source exits, handle_watch_result enqueues a WatchEvent
941        // with exited=true and unregisters the watcher. wait_for_event must still
942        // return that event even though no watchers remain active.
943        let hub = Arc::new(WatchHub::new());
944        let wid = hub.register(
945            WatchSource::Agent {
946                handle: "a1".into(),
947            },
948            "done".into(),
949            WatchMode::Once,
950            Duration::from_secs(10),
951        );
952        // Simulate source-exited: enqueue event then unregister (mirrors handle_watch_result)
953        hub.enqueue_event(WatchEvent {
954            watcher_id: wid.clone(),
955            source: WatchSource::Agent {
956                handle: "a1".into(),
957            },
958            pattern: "done".into(),
959            row: None,
960            col: None,
961            text: String::new(),
962            timestamp: chrono::Utc::now(),
963            timed_out: false,
964            exited: true,
965            timeout: Duration::from_secs(10),
966        });
967        hub.unregister(&wid);
968        assert!(!hub.has_active_watchers());
969        // Must still return the pending exited event despite no active watchers
970        let evt = hub.wait_for_event(Duration::from_millis(100)).await;
971        assert!(
972            evt.is_some(),
973            "exited event must be consumable via wait_for_event"
974        );
975        let evt = evt.unwrap();
976        assert!(evt.exited);
977        assert_eq!(evt.pattern, "done");
978    }
979
980    #[test]
981    fn format_event_text_exited_includes_already_exited() {
982        let evt = WatchEvent {
983            watcher_id: "w_xyz".into(),
984            source: WatchSource::Agent {
985                handle: "a1".into(),
986            },
987            pattern: "done".into(),
988            row: None,
989            col: None,
990            text: String::new(),
991            timestamp: chrono::Utc::now(),
992            timed_out: false,
993            exited: true,
994            timeout: Duration::from_secs(120),
995        };
996        let text = format_watch_event_text(&evt);
997        assert!(text.contains("already exited"));
998        assert!(text.contains("a1"));
999        assert!(text.contains("done"));
1000    }
1001
1002    struct ScriptedWatchable {
1003        results: Vec<WatchResult>,
1004        idx: Mutex<usize>,
1005    }
1006
1007    impl Watchable for ScriptedWatchable {
1008        fn watch_output(
1009            self: Arc<Self>,
1010            _pattern: String,
1011            _cancel: CancellationToken,
1012        ) -> std::pin::Pin<Box<dyn std::future::Future<Output = WatchResult> + Send>> {
1013            let mut idx = self.idx.lock().unwrap();
1014            let i = *idx;
1015            *idx += 1;
1016            let result = self
1017                .results
1018                .get(i)
1019                .cloned()
1020                .unwrap_or(WatchResult::SourceExited);
1021            Box::pin(async move {
1022                tokio::time::sleep(Duration::from_millis(10)).await;
1023                result
1024            })
1025        }
1026    }
1027
1028    #[tokio::test]
1029    async fn persist_watcher_fires_on_every_match_then_cleans_up() {
1030        let hub = Arc::new(WatchHub::new());
1031        let wid = hub.register(
1032            WatchSource::Bash {
1033                handle: "b1".into(),
1034            },
1035            "done".into(),
1036            WatchMode::Persist,
1037            Duration::from_secs(10),
1038        );
1039        let watchable: Arc<dyn Watchable> = Arc::new(ScriptedWatchable {
1040            results: vec![
1041                WatchResult::Matched {
1042                    row: None,
1043                    col: Some(0),
1044                    text: "done".into(),
1045                },
1046                WatchResult::Matched {
1047                    row: None,
1048                    col: Some(4),
1049                    text: "done".into(),
1050                },
1051                WatchResult::SourceExited,
1052            ],
1053            idx: Mutex::new(0),
1054        });
1055        spawn_watcher(
1056            Arc::clone(&hub),
1057            wid,
1058            WatchSource::Bash {
1059                handle: "b1".into(),
1060            },
1061            "done".into(),
1062            WatchMode::Persist,
1063            Duration::from_secs(10),
1064            watchable,
1065        );
1066
1067        let e1 = hub
1068            .wait_for_event(Duration::from_secs(2))
1069            .await
1070            .expect("first match event");
1071        assert!(!e1.exited && !e1.timed_out, "first event should be a match");
1072        let e2 = hub
1073            .wait_for_event(Duration::from_secs(2))
1074            .await
1075            .expect("second match event");
1076        assert!(
1077            !e2.exited && !e2.timed_out,
1078            "second event should be a match"
1079        );
1080        let e3 = hub
1081            .wait_for_event(Duration::from_secs(2))
1082            .await
1083            .expect("exited event");
1084        assert!(e3.exited, "third event should be source-exited");
1085
1086        tokio::time::sleep(Duration::from_millis(50)).await;
1087        assert!(
1088            !hub.has_active_watchers(),
1089            "persist watcher must be removed after source exits"
1090        );
1091    }
1092
1093    #[tokio::test]
1094    async fn persist_watcher_timeout_removes_watcher() {
1095        struct NeverMatch;
1096        impl Watchable for NeverMatch {
1097            fn watch_output(
1098                self: Arc<Self>,
1099                _pattern: String,
1100                cancel: CancellationToken,
1101            ) -> std::pin::Pin<Box<dyn std::future::Future<Output = WatchResult> + Send>>
1102            {
1103                Box::pin(async move {
1104                    cancel.cancelled().await;
1105                    WatchResult::Cancelled
1106                })
1107            }
1108        }
1109
1110        let hub = Arc::new(WatchHub::new());
1111        let wid = hub.register(
1112            WatchSource::Bash {
1113                handle: "b1".into(),
1114            },
1115            "done".into(),
1116            WatchMode::Persist,
1117            Duration::from_millis(50),
1118        );
1119        spawn_watcher(
1120            Arc::clone(&hub),
1121            wid,
1122            WatchSource::Bash {
1123                handle: "b1".into(),
1124            },
1125            "done".into(),
1126            WatchMode::Persist,
1127            Duration::from_millis(50),
1128            Arc::new(NeverMatch) as Arc<dyn Watchable>,
1129        );
1130
1131        let evt = hub
1132            .wait_for_event(Duration::from_secs(2))
1133            .await
1134            .expect("timeout event");
1135        assert!(evt.timed_out, "should receive a timeout event");
1136
1137        tokio::time::sleep(Duration::from_millis(50)).await;
1138        assert!(
1139            !hub.has_active_watchers(),
1140            "persist watcher must be removed after timeout, not linger forever"
1141        );
1142    }
1143
1144    #[tokio::test]
1145    async fn persist_watcher_stops_on_cancel() {
1146        struct NeverMatch;
1147        impl Watchable for NeverMatch {
1148            fn watch_output(
1149                self: Arc<Self>,
1150                _pattern: String,
1151                cancel: CancellationToken,
1152            ) -> std::pin::Pin<Box<dyn std::future::Future<Output = WatchResult> + Send>>
1153            {
1154                Box::pin(async move {
1155                    cancel.cancelled().await;
1156                    WatchResult::Cancelled
1157                })
1158            }
1159        }
1160
1161        let hub = Arc::new(WatchHub::new());
1162        let wid = hub.register(
1163            WatchSource::Bash {
1164                handle: "b1".into(),
1165            },
1166            "done".into(),
1167            WatchMode::Persist,
1168            Duration::from_secs(10),
1169        );
1170        spawn_watcher(
1171            Arc::clone(&hub),
1172            wid.clone(),
1173            WatchSource::Bash {
1174                handle: "b1".into(),
1175            },
1176            "done".into(),
1177            WatchMode::Persist,
1178            Duration::from_secs(10),
1179            Arc::new(NeverMatch) as Arc<dyn Watchable>,
1180        );
1181
1182        hub.unregister(&wid);
1183
1184        tokio::time::sleep(Duration::from_millis(100)).await;
1185        assert!(
1186            !hub.has_active_watchers(),
1187            "persist watcher must stop and be removed when cancelled"
1188        );
1189    }
1190}
atman_runtime/watch.rs

atman_runtime/
watch.rs
