basis/run/
sink.rs

1//! Where a run's events go.
2//!
3//! `basis spawn --json` writes JSONL to stdout; an in-process host wants a
4//! callback or a buffer; P2's ACP server will want a protocol notification.
5//! All of them are the same stream, so they are all the same trait.
6//!
7//! A host running many runs at once wants one view of all of them; that is
8//! [`fan_in`](mod@fan_in), which is the same trait again with a tag on the
9//! front.
10
11mod fan_in;
12
13use std::io::Write;
14
15use crate::event::{Event, JsonlWriter};
16
17pub use fan_in::{EventFanIn, MergedEvents, TaggedEvent, TaggedSink};
18
19/// A destination for run events.
20///
21/// Emission happens on a background task, so a sink must be `Send`. Returning
22/// an error stops emission for the rest of the run — a client that has gone
23/// away should not cost the run a full transcript of failed writes.
24///
25/// The run itself carries on. That task is also the one answering approval
26/// requests, and mentra blocks the turn until one is answered, so giving up on
27/// it would turn a broken pipe into a hung agent.
28pub trait EventSink: Send + 'static {
29    fn emit(&mut self, event: Event) -> std::io::Result<()>;
30}
31
32/// A boxed sink is a sink, so a host can choose one at runtime — `--json` or a
33/// progress pane or nothing — and still satisfy the `S: EventSink` that
34/// [`execute_with_approver`](super::PreparedRun::execute_with_approver) asks
35/// for. Without this the choice
36/// has to be made in the type system, which is to say at every call site that
37/// forwards it.
38impl EventSink for Box<dyn EventSink> {
39    fn emit(&mut self, event: Event) -> std::io::Result<()> {
40        (**self).emit(event)
41    }
42}
43
44impl<W: Write + Send + 'static> EventSink for JsonlWriter<W> {
45    fn emit(&mut self, event: Event) -> std::io::Result<()> {
46        self.write(event).map(|_| ())
47    }
48}
49
50/// Keeps every event in memory. The natural sink for tests, and for a host
51/// that wants the whole run before deciding what to do with it.
52#[derive(Debug, Default, Clone, PartialEq)]
53pub struct CollectingSink {
54    events: Vec<Event>,
55}
56
57impl CollectingSink {
58    pub fn new() -> Self {
59        Self::default()
60    }
61
62    pub fn events(&self) -> &[Event] {
63        &self.events
64    }
65
66    pub fn into_events(self) -> Vec<Event> {
67        self.events
68    }
69}
70
71impl EventSink for CollectingSink {
72    fn emit(&mut self, event: Event) -> std::io::Result<()> {
73        self.events.push(event);
74        Ok(())
75    }
76}
77
78/// Discards every event, for a host that only wants the final message.
79#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
80pub struct NullSink;
81
82impl EventSink for NullSink {
83    fn emit(&mut self, _event: Event) -> std::io::Result<()> {
84        Ok(())
85    }
86}
87
88/// Calls a closure per event.
89pub struct FnSink<F>(F)
90where
91    F: FnMut(Event) -> std::io::Result<()> + Send + 'static;
92
93impl<F> FnSink<F>
94where
95    F: FnMut(Event) -> std::io::Result<()> + Send + 'static,
96{
97    pub fn new(callback: F) -> Self {
98        Self(callback)
99    }
100}
101
102impl<F> EventSink for FnSink<F>
103where
104    F: FnMut(Event) -> std::io::Result<()> + Send + 'static,
105{
106    fn emit(&mut self, event: Event) -> std::io::Result<()> {
107        (self.0)(event)
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114    use crate::event::RunOutcome;
115
116    fn delta(text: &str) -> Event {
117        Event::AssistantDelta {
118            text: text.to_string(),
119        }
120    }
121
122    #[test]
123    fn collecting_sink_keeps_order() {
124        let mut sink = CollectingSink::new();
125        sink.emit(delta("a")).expect("emits");
126        sink.emit(delta("b")).expect("emits");
127
128        assert_eq!(sink.into_events(), vec![delta("a"), delta("b")]);
129    }
130
131    #[test]
132    fn null_sink_accepts_everything() {
133        let mut sink = NullSink;
134
135        assert!(
136            sink.emit(Event::RunFinished {
137                outcome: RunOutcome::Ok,
138                stopped_by: None,
139                usage: None
140            })
141            .is_ok()
142        );
143    }
144
145    #[test]
146    fn fn_sink_forwards_to_the_closure() {
147        let (tx, rx) = std::sync::mpsc::channel();
148        let mut sink = FnSink::new(move |event| {
149            tx.send(event).expect("receiver alive");
150            Ok(())
151        });
152
153        sink.emit(delta("x")).expect("emits");
154
155        assert_eq!(rx.recv().expect("an event"), delta("x"));
156    }
157
158    #[test]
159    fn a_jsonl_writer_is_a_sink() {
160        let mut sink = JsonlWriter::new(Vec::new());
161        sink.emit(delta("hi")).expect("emits");
162
163        let written = String::from_utf8(sink.into_inner()).expect("utf-8");
164        assert!(written.contains("\"type\":\"assistant_delta\""));
165    }
166
167    #[test]
168    fn a_sink_chosen_at_runtime_is_still_a_sink() {
169        // What a host with a `--json` flag actually has: one variable, two
170        // possible destinations, and a run that only knows it takes a sink.
171        let mut sink: Box<dyn EventSink> = if cfg!(test) {
172            Box::new(CollectingSink::new())
173        } else {
174            Box::new(NullSink)
175        };
176
177        sink.emit(delta("boxed")).expect("emits");
178    }
179}
basis/run/sink.rs

basis/run/
sink.rs
