rightkit_agent_engine_protocol/
tool_output.rs1use std::fmt;
7use std::sync::Arc;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum ToolOutputStream {
12 Stdout,
13 Stderr,
14}
15
16impl ToolOutputStream {
17 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#[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}