vv_agent/
event_store.rs

1use std::fmt;
2use std::fs::{File, OpenOptions};
3use std::io::{BufRead, BufReader, Write};
4use std::path::{Path, PathBuf};
5
6use crate::checkpoint::EventCursor;
7use crate::events::RunEvent;
8
9pub trait RunEventStore: Send + Sync {
10    fn append(&self, event: &RunEvent) -> Result<(), EventStoreError>;
11    fn replay(&self, query: RunEventReplayQuery) -> Result<RunEventIter, EventStoreError>;
12
13    fn append_once(
14        &self,
15        _event_id: &str,
16        _payload_digest: &str,
17        _event: &RunEvent,
18    ) -> Result<Option<EventCursor>, EventStoreError> {
19        Ok(None)
20    }
21}
22
23pub type RunEventIter = Box<dyn Iterator<Item = Result<RunEvent, EventStoreError>> + Send>;
24
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct RunEventReplayQuery {
27    run_id: String,
28    include_children: bool,
29}
30
31impl RunEventReplayQuery {
32    pub fn run(run_id: impl Into<String>) -> Self {
33        Self {
34            run_id: run_id.into(),
35            include_children: true,
36        }
37    }
38
39    pub fn include_children(mut self, include_children: bool) -> Self {
40        self.include_children = include_children;
41        self
42    }
43
44    pub fn run_id(&self) -> &str {
45        &self.run_id
46    }
47
48    pub fn should_include_children(&self) -> bool {
49        self.include_children
50    }
51}
52
53#[derive(Debug, Clone)]
54pub struct JsonlRunEventStore {
55    path: PathBuf,
56}
57
58impl JsonlRunEventStore {
59    pub fn new(path: impl AsRef<Path>) -> Self {
60        Self {
61            path: path.as_ref().to_path_buf(),
62        }
63    }
64}
65
66impl RunEventStore for JsonlRunEventStore {
67    fn append(&self, event: &RunEvent) -> Result<(), EventStoreError> {
68        if let Some(parent) = self.path.parent() {
69            std::fs::create_dir_all(parent).map_err(EventStoreError::io)?;
70        }
71        let mut file = OpenOptions::new()
72            .create(true)
73            .append(true)
74            .open(&self.path)
75            .map_err(EventStoreError::io)?;
76        let line = serde_json::to_string(event).map_err(EventStoreError::json)?;
77        writeln!(file, "{line}").map_err(EventStoreError::io)?;
78        Ok(())
79    }
80
81    fn replay(&self, query: RunEventReplayQuery) -> Result<RunEventIter, EventStoreError> {
82        let file = match File::open(&self.path) {
83            Ok(file) => file,
84            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
85                return Ok(Box::new(std::iter::empty()));
86            }
87            Err(error) => return Err(EventStoreError::io(error)),
88        };
89        let mut lines = BufReader::new(file).lines().enumerate();
90        let mut stopped = false;
91
92        Ok(Box::new(std::iter::from_fn(move || {
93            if stopped {
94                return None;
95            }
96
97            loop {
98                let Some((line_index, line)) = lines.next() else {
99                    stopped = true;
100                    return None;
101                };
102                let line_number = line_index + 1;
103                let line = match line {
104                    Ok(line) => line,
105                    Err(error) => {
106                        stopped = true;
107                        return Some(Err(EventStoreError::io(error)));
108                    }
109                };
110                let event: RunEvent = match serde_json::from_str(&line) {
111                    Ok(event) => event,
112                    Err(_) => {
113                        stopped = true;
114                        return Some(Err(EventStoreError::corrupt_line(line_number)));
115                    }
116                };
117                let include = event.run_id() == query.run_id()
118                    || (query.should_include_children()
119                        && event.parent_run_id() == Some(query.run_id()));
120                if include {
121                    return Some(Ok(event));
122                }
123            }
124        })))
125    }
126}
127
128#[derive(Debug, Clone, PartialEq, Eq)]
129pub struct EventStoreError {
130    code: &'static str,
131    message: String,
132    line_number: Option<usize>,
133}
134
135impl EventStoreError {
136    pub fn new(code: &'static str, message: impl Into<String>) -> Self {
137        Self {
138            code,
139            message: message.into(),
140            line_number: None,
141        }
142    }
143
144    fn io(error: std::io::Error) -> Self {
145        Self {
146            code: "event_store_io_error",
147            message: error.to_string(),
148            line_number: None,
149        }
150    }
151
152    fn json(error: serde_json::Error) -> Self {
153        Self {
154            code: "event_store_serialization_error",
155            message: error.to_string(),
156            line_number: None,
157        }
158    }
159
160    fn corrupt_line(line_number: usize) -> Self {
161        Self {
162            code: "event_store_corrupt_line",
163            message: format!("event store corrupt line {line_number}"),
164            line_number: Some(line_number),
165        }
166    }
167
168    pub fn code(&self) -> &'static str {
169        self.code
170    }
171
172    pub fn line_number(&self) -> Option<usize> {
173        self.line_number
174    }
175}
176
177impl fmt::Display for EventStoreError {
178    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
179        formatter.write_str(&self.message)
180    }
181}
182
183impl std::error::Error for EventStoreError {}
vv_agent/event_store.rs

vv_agent/
event_store.rs
