use std::io::Write;
use super::{Event, EventLine};
#[derive(Debug)]
pub struct JsonlWriter<W: Write> {
writer: W,
next_seq: u64,
}
impl<W: Write> JsonlWriter<W> {
pub fn new(writer: W) -> Self {
Self {
writer,
next_seq: 0,
}
}
pub fn write(&mut self, event: Event) -> std::io::Result<u64> {
let seq = self.next_seq;
let line = EventLine::new(seq, event);
let encoded = serde_json::to_string(&line).map_err(std::io::Error::other)?;
writeln!(self.writer, "{encoded}")?;
self.writer.flush()?;
self.next_seq += 1;
Ok(seq)
}
pub fn next_seq(&self) -> u64 {
self.next_seq
}
pub fn into_inner(self) -> W {
self.writer
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::event::RunOutcome;
fn lines(buffer: &[u8]) -> Vec<serde_json::Value> {
String::from_utf8(buffer.to_vec())
.expect("utf-8")
.lines()
.map(|line| serde_json::from_str(line).expect("each line is json"))
.collect()
}
#[test]
fn every_event_gets_the_next_sequence_number() {
let mut writer = JsonlWriter::new(Vec::new());
writer
.write(Event::AssistantDelta {
text: "a".to_string(),
})
.expect("writes");
writer
.write(Event::AssistantDelta {
text: "b".to_string(),
})
.expect("writes");
let written = lines(&writer.into_inner());
assert_eq!(written[0]["seq"], 0);
assert_eq!(written[1]["seq"], 1);
}
#[test]
fn each_event_is_exactly_one_line() {
let mut writer = JsonlWriter::new(Vec::new());
writer
.write(Event::AssistantMessage {
text: "one\ntwo\nthree".to_string(),
})
.expect("writes");
writer
.write(Event::RunFinished {
outcome: RunOutcome::Ok,
stopped_by: None,
})
.expect("writes");
let buffer = writer.into_inner();
let text = String::from_utf8(buffer.clone()).expect("utf-8");
assert_eq!(text.lines().count(), 2);
let written = lines(&buffer);
assert_eq!(written[0]["text"], "one\ntwo\nthree");
}
#[test]
fn next_seq_reports_what_the_next_write_will_use() {
let mut writer = JsonlWriter::new(Vec::new());
assert_eq!(writer.next_seq(), 0);
writer
.write(Event::CompactionStarted {
agent_id: "a1".to_string(),
})
.expect("writes");
assert_eq!(writer.next_seq(), 1);
}
#[test]
fn a_broken_pipe_surfaces_rather_than_being_swallowed() {
struct Broken;
impl Write for Broken {
fn write(&mut self, _buf: &[u8]) -> std::io::Result<usize> {
Err(std::io::Error::new(std::io::ErrorKind::BrokenPipe, "gone"))
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
let mut writer = JsonlWriter::new(Broken);
let error = writer
.write(Event::AssistantDelta {
text: "x".to_string(),
})
.expect_err("write fails");
assert_eq!(error.kind(), std::io::ErrorKind::BrokenPipe);
}
}