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`](super::PreparedRun::execute) asks for. Without this the choice
35/// has to be made in the type system, which is to say at every call site that
36/// forwards it.
37impl EventSink for Box<dyn EventSink> {
38    fn emit(&mut self, event: Event) -> std::io::Result<()> {
39        (**self).emit(event)
40    }
41}
42
43impl<W: Write + Send + 'static> EventSink for JsonlWriter<W> {
44    fn emit(&mut self, event: Event) -> std::io::Result<()> {
45        self.write(event).map(|_| ())
46    }
47}
48
49/// Keeps every event in memory. The natural sink for tests, and for a host
50/// that wants the whole run before deciding what to do with it.
51#[derive(Debug, Default, Clone, PartialEq)]
52pub struct CollectingSink {
53    events: Vec<Event>,
54}
55
56impl CollectingSink {
57    pub fn new() -> Self {
58        Self::default()
59    }
60
61    pub fn events(&self) -> &[Event] {
62        &self.events
63    }
64
65    pub fn into_events(self) -> Vec<Event> {
66        self.events
67    }
68}
69
70impl EventSink for CollectingSink {
71    fn emit(&mut self, event: Event) -> std::io::Result<()> {
72        self.events.push(event);
73        Ok(())
74    }
75}
76
77/// Discards every event, for a host that only wants the final message.
78#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
79pub struct NullSink;
80
81impl EventSink for NullSink {
82    fn emit(&mut self, _event: Event) -> std::io::Result<()> {
83        Ok(())
84    }
85}
86
87/// Calls a closure per event.
88pub struct FnSink<F>(F)
89where
90    F: FnMut(Event) -> std::io::Result<()> + Send + 'static;
91
92impl<F> FnSink<F>
93where
94    F: FnMut(Event) -> std::io::Result<()> + Send + 'static,
95{
96    pub fn new(callback: F) -> Self {
97        Self(callback)
98    }
99}
100
101impl<F> EventSink for FnSink<F>
102where
103    F: FnMut(Event) -> std::io::Result<()> + Send + 'static,
104{
105    fn emit(&mut self, event: Event) -> std::io::Result<()> {
106        (self.0)(event)
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113    use crate::event::RunOutcome;
114
115    fn delta(text: &str) -> Event {
116        Event::AssistantDelta {
117            text: text.to_string(),
118        }
119    }
120
121    #[test]
122    fn collecting_sink_keeps_order() {
123        let mut sink = CollectingSink::new();
124        sink.emit(delta("a")).expect("emits");
125        sink.emit(delta("b")).expect("emits");
126
127        assert_eq!(sink.into_events(), vec![delta("a"), delta("b")]);
128    }
129
130    #[test]
131    fn null_sink_accepts_everything() {
132        let mut sink = NullSink;
133
134        assert!(
135            sink.emit(Event::RunFinished {
136                outcome: RunOutcome::Ok,
137                stopped_by: None
138            })
139            .is_ok()
140        );
141    }
142
143    #[test]
144    fn fn_sink_forwards_to_the_closure() {
145        let (tx, rx) = std::sync::mpsc::channel();
146        let mut sink = FnSink::new(move |event| {
147            tx.send(event).expect("receiver alive");
148            Ok(())
149        });
150
151        sink.emit(delta("x")).expect("emits");
152
153        assert_eq!(rx.recv().expect("an event"), delta("x"));
154    }
155
156    #[test]
157    fn a_jsonl_writer_is_a_sink() {
158        let mut sink = JsonlWriter::new(Vec::new());
159        sink.emit(delta("hi")).expect("emits");
160
161        let written = String::from_utf8(sink.into_inner()).expect("utf-8");
162        assert!(written.contains("\"type\":\"assistant_delta\""));
163    }
164
165    #[test]
166    fn a_sink_chosen_at_runtime_is_still_a_sink() {
167        // What a host with a `--json` flag actually has: one variable, two
168        // possible destinations, and a run that only knows it takes a sink.
169        let mut sink: Box<dyn EventSink> = if cfg!(test) {
170            Box::new(CollectingSink::new())
171        } else {
172            Box::new(NullSink)
173        };
174
175        sink.emit(delta("boxed")).expect("emits");
176    }
177}
basis/run/sink.rs

basis/run/
sink.rs
