Skip to main content

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::events::RunEvent;
7
8pub trait RunEventStore: Send + Sync {
9    fn append(&self, event: &RunEvent) -> Result<(), EventStoreError>;
10    fn replay(&self, query: RunEventReplayQuery) -> Result<RunEventIter, EventStoreError>;
11}
12
13pub type RunEventIter = Box<dyn Iterator<Item = Result<RunEvent, EventStoreError>> + Send>;
14
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct RunEventReplayQuery {
17    run_id: String,
18    include_children: bool,
19}
20
21impl RunEventReplayQuery {
22    pub fn run(run_id: impl Into<String>) -> Self {
23        Self {
24            run_id: run_id.into(),
25            include_children: false,
26        }
27    }
28
29    pub fn include_children(mut self, include_children: bool) -> Self {
30        self.include_children = include_children;
31        self
32    }
33
34    pub fn run_id(&self) -> &str {
35        &self.run_id
36    }
37
38    pub fn should_include_children(&self) -> bool {
39        self.include_children
40    }
41}
42
43#[derive(Debug, Clone)]
44pub struct JsonlRunEventStore {
45    path: PathBuf,
46}
47
48impl JsonlRunEventStore {
49    pub fn new(path: impl AsRef<Path>) -> Self {
50        Self {
51            path: path.as_ref().to_path_buf(),
52        }
53    }
54}
55
56impl RunEventStore for JsonlRunEventStore {
57    fn append(&self, event: &RunEvent) -> Result<(), EventStoreError> {
58        if let Some(parent) = self.path.parent() {
59            std::fs::create_dir_all(parent).map_err(EventStoreError::io)?;
60        }
61        let mut file = OpenOptions::new()
62            .create(true)
63            .append(true)
64            .open(&self.path)
65            .map_err(EventStoreError::io)?;
66        let line = serde_json::to_string(event).map_err(EventStoreError::json)?;
67        writeln!(file, "{line}").map_err(EventStoreError::io)?;
68        Ok(())
69    }
70
71    fn replay(&self, query: RunEventReplayQuery) -> Result<RunEventIter, EventStoreError> {
72        if !self.path.exists() {
73            return Ok(Box::new(std::iter::empty()));
74        }
75
76        let file = File::open(&self.path).map_err(EventStoreError::io)?;
77        let reader = BufReader::new(file);
78        let mut events = Vec::new();
79        for line in reader.lines() {
80            let line = line.map_err(EventStoreError::io)?;
81            let event: RunEvent = serde_json::from_str(&line).map_err(EventStoreError::json)?;
82            let include = event.run_id() == query.run_id()
83                || (query.should_include_children()
84                    && event.parent_run_id() == Some(query.run_id()));
85            if include {
86                events.push(event);
87            }
88        }
89
90        Ok(Box::new(events.into_iter().map(Ok)))
91    }
92}
93
94#[derive(Debug, Clone, PartialEq, Eq)]
95pub struct EventStoreError {
96    message: String,
97}
98
99impl EventStoreError {
100    fn io(error: std::io::Error) -> Self {
101        Self {
102            message: error.to_string(),
103        }
104    }
105
106    fn json(error: serde_json::Error) -> Self {
107        Self {
108            message: error.to_string(),
109        }
110    }
111}
112
113impl fmt::Display for EventStoreError {
114    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
115        formatter.write_str(&self.message)
116    }
117}
118
119impl std::error::Error for EventStoreError {}