@ondeinference / ed
Code Commits Issues Pulls Cloud Agent CI/CD
release/1.3.1
main release/1.3.1
rs 174 lines 5.77 KB
Raw
1 //! Durable sessions: one JSON file per session under `<data_dir>/sessions/`.
2 //!
3 //! The history stored here is the OpenAI-style message list without the system prompt
4 //! (that is rebuilt every turn). Tool results are `role: "tool"` messages, which is what
5 //! lets `session/load` replay tool calls with their output.
6
7 use std::path::{Path, PathBuf};
8
9 use anyhow::{Context, Result, bail};
10 use serde::{Deserialize, Serialize};
11 use serde_json::Value;
12
13 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
14 pub struct StoredSession {
15 pub cwd: PathBuf,
16 #[serde(default)]
17 pub roots: Vec<PathBuf>,
18 pub model: String,
19 #[serde(default)]
20 pub title: Option<String>,
21 pub updated_at: u64,
22 #[serde(default)]
23 pub messages: Vec<Value>,
24 }
25
26 #[derive(Debug, Clone)]
27 pub struct SessionStore {
28 dir: PathBuf,
29 }
30
31 /// Session ids end up in file names: allow only characters that cannot escape the directory.
32 pub fn sanitize_id(id: &str) -> Result<&str> {
33 let ok = !id.is_empty()
34 && id.len() <= 128
35 && id
36 .chars()
37 .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_');
38 if !ok {
39 bail!("invalid session id");
40 }
41 Ok(id)
42 }
43
44 impl SessionStore {
45 pub fn new(data_dir: &Path) -> Self {
46 Self {
47 dir: data_dir.join("sessions"),
48 }
49 }
50
51 fn path(&self, id: &str) -> Result<PathBuf> {
52 Ok(self.dir.join(format!("{}.json", sanitize_id(id)?)))
53 }
54
55 /// Atomic write: temp file in the same directory, then rename.
56 pub fn save(&self, id: &str, session: &StoredSession) -> Result<()> {
57 let path = self.path(id)?;
58 std::fs::create_dir_all(&self.dir)
59 .with_context(|| format!("creating {}", self.dir.display()))?;
60 let tmp = self.dir.join(format!(".{id}.{}.tmp", uuid::Uuid::new_v4()));
61 std::fs::write(&tmp, serde_json::to_vec(session)?)?;
62 std::fs::rename(&tmp, &path).inspect_err(|_| {
63 let _ = std::fs::remove_file(&tmp);
64 })?;
65 Ok(())
66 }
67
68 /// `Ok(None)` when no such session exists (or the id is not a valid one).
69 pub fn load(&self, id: &str) -> Result<Option<StoredSession>> {
70 let Ok(path) = self.path(id) else {
71 return Ok(None);
72 };
73 match std::fs::read(&path) {
74 Ok(bytes) => Ok(Some(
75 serde_json::from_slice(&bytes)
76 .with_context(|| format!("parsing {}", path.display()))?,
77 )),
78 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
79 Err(e) => Err(e).with_context(|| format!("reading {}", path.display())),
80 }
81 }
82
83 /// All stored sessions, newest first. Unreadable files are skipped.
84 pub fn list(&self) -> Vec<(String, StoredSession)> {
85 let Ok(rd) = std::fs::read_dir(&self.dir) else {
86 return Vec::new();
87 };
88 let mut out: Vec<(String, StoredSession)> = rd
89 .filter_map(|e| {
90 let path = e.ok()?.path();
91 if path.extension()? != "json" {
92 return None;
93 }
94 let id = path.file_stem()?.to_str()?.to_string();
95 let session = serde_json::from_slice(&std::fs::read(&path).ok()?).ok()?;
96 Some((id, session))
97 })
98 .collect();
99 out.sort_by(|a, b| b.1.updated_at.cmp(&a.1.updated_at).then(a.0.cmp(&b.0)));
100 out
101 }
102
103 /// Returns whether a file was removed.
104 pub fn delete(&self, id: &str) -> Result<bool> {
105 let Ok(path) = self.path(id) else {
106 return Ok(false);
107 };
108 match std::fs::remove_file(&path) {
109 Ok(()) => Ok(true),
110 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(false),
111 Err(e) => Err(e).with_context(|| format!("deleting {}", path.display())),
112 }
113 }
114 }
115
116 #[cfg(test)]
117 mod tests {
118 use super::*;
119 use serde_json::json;
120
121 fn tmp() -> PathBuf {
122 let d = std::env::temp_dir().join(format!("sf-store-{}", uuid::Uuid::new_v4()));
123 std::fs::create_dir_all(&d).unwrap();
124 d
125 }
126
127 fn sample(updated_at: u64) -> StoredSession {
128 StoredSession {
129 cwd: "/work".into(),
130 roots: vec!["/other".into()],
131 model: "m".into(),
132 title: Some("hello".into()),
133 updated_at,
134 messages: vec![
135 json!({"role": "user", "content": "hi"}),
136 json!({"role": "tool", "tool_call_id": "c1", "content": "out"}),
137 ],
138 }
139 }
140
141 #[test]
142 fn round_trip_list_delete() {
143 let dir = tmp();
144 let store = SessionStore::new(&dir);
145 assert!(store.load("a").unwrap().is_none());
146 store.save("a", &sample(10)).unwrap();
147 store.save("b", &sample(20)).unwrap();
148 assert_eq!(store.load("a").unwrap().unwrap(), sample(10));
149 let ids: Vec<_> = store.list().into_iter().map(|(id, _)| id).collect();
150 assert_eq!(ids, ["b", "a"]);
151 assert!(store.delete("a").unwrap());
152 assert!(!store.delete("a").unwrap());
153 assert!(store.load("a").unwrap().is_none());
154 // No temp files left behind.
155 let leftovers = std::fs::read_dir(dir.join("sessions"))
156 .unwrap()
157 .filter(|e| e.as_ref().unwrap().path().extension().unwrap() != "json")
158 .count();
159 assert_eq!(leftovers, 0);
160 std::fs::remove_dir_all(dir).ok();
161 }
162
163 #[test]
164 fn rejects_path_traversal_ids() {
165 let dir = tmp();
166 let store = SessionStore::new(&dir);
167 for bad in ["../x", "a/b", "", "a.b", "..", "a\\b"] {
168 assert!(store.save(bad, &sample(1)).is_err(), "{bad}");
169 assert!(store.load(bad).unwrap().is_none());
170 assert!(!store.delete(bad).unwrap());
171 }
172 std::fs::remove_dir_all(dir).ok();
173 }
174 }