Skip to main content

rightkit_agent_engine_protocol/
tool_output.rs

1//! Live tool output hook: a plain-data sink so a long-running process tool
2//! (shell/exec) can report output chunks while it runs without owning any
3//! session state. The engine installs a sink per tool call and turns the
4//! (coalesced) chunks into `toolOutputDelta` protocol events.
5
6use std::fmt;
7use std::sync::Arc;
8
9/// Which pipe a chunk was read from.
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum ToolOutputStream {
12    Stdout,
13    Stderr,
14}
15
16impl ToolOutputStream {
17    /// Wire spelling used by `toolOutputDelta.stream`.
18    pub fn as_str(self) -> &'static str {
19        match self {
20            Self::Stdout => "stdout",
21            Self::Stderr => "stderr",
22        }
23    }
24}
25
26type OutputFn = dyn Fn(ToolOutputStream, &str) + Send + Sync;
27
28/// Cloneable, optional output sink. Default is a no-op, so a tool can call
29/// [`ToolOutputNotifier::notify`] unconditionally.
30#[derive(Clone, Default)]
31pub struct ToolOutputNotifier(Option<Arc<OutputFn>>);
32
33impl ToolOutputNotifier {
34    pub fn new(f: impl Fn(ToolOutputStream, &str) + Send + Sync + 'static) -> Self {
35        Self(Some(Arc::new(f)))
36    }
37
38    pub fn is_set(&self) -> bool {
39        self.0.is_some()
40    }
41
42    pub fn notify(&self, stream: ToolOutputStream, text: &str) {
43        if let Some(f) = &self.0 {
44            f(stream, text);
45        }
46    }
47}
48
49impl fmt::Debug for ToolOutputNotifier {
50    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
51        f.write_str(if self.0.is_some() {
52            "ToolOutputNotifier(set)"
53        } else {
54            "ToolOutputNotifier(none)"
55        })
56    }
57}
58
59impl PartialEq for ToolOutputNotifier {
60    fn eq(&self, other: &Self) -> bool {
61        match (&self.0, &other.0) {
62            (None, None) => true,
63            (Some(a), Some(b)) => Arc::ptr_eq(a, b),
64            _ => false,
65        }
66    }
67}
68
69impl Eq for ToolOutputNotifier {}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use std::sync::Mutex;
75
76    #[test]
77    fn notifier_delivers_and_default_is_silent() {
78        ToolOutputNotifier::default().notify(ToolOutputStream::Stdout, "ignored");
79        let seen = Arc::new(Mutex::new(Vec::new()));
80        let sink = seen.clone();
81        let notifier = ToolOutputNotifier::new(move |stream, text| {
82            sink.lock()
83                .unwrap()
84                .push((stream.as_str(), text.to_string()));
85        });
86        notifier.notify(ToolOutputStream::Stderr, "warn");
87        assert!(notifier.is_set());
88        assert_eq!(*seen.lock().unwrap(), vec![("stderr", "warn".to_string())]);
89    }
90}