use std::io;
use std::os::unix::fs::PermissionsExt as _;
use std::path::{Path, PathBuf};
use std::time::{SystemTime, UNIX_EPOCH};
use qframe::storage::atomic_write;
use crate::conversation::Conversation;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Store {
file: PathBuf,
unreadable: bool,
}
impl Store {
#[must_use]
pub fn new(root: &Path, folder: &Path) -> Self {
let name = folder.file_name().map_or_else(|| "root".to_owned(), |name| name.to_string_lossy().into_owned());
let name: String =
name.chars().map(|c| if c.is_alphanumeric() || c == '-' || c == '_' { c } else { '_' }).collect();
let file = root.join(format!("{name}-{:016x}.json", fnv(folder.as_os_str().as_encoded_bytes())));
Self { file, unreadable: false }
}
#[must_use]
pub fn file(&self) -> &Path {
&self.file
}
pub fn load(&mut self) -> Conversation {
let Ok(text) = std::fs::read_to_string(&self.file) else {
self.unreadable = std::fs::symlink_metadata(&self.file).is_ok();
return Conversation::default();
};
let conversation =
serde_json::from_str(&text).ok().filter(Conversation::holds).map(|value| Conversation::from_json(&value));
self.unreadable = conversation.is_none();
conversation.unwrap_or_default()
}
pub fn save(&mut self, conversation: &Conversation) -> io::Result<()> {
if self.unreadable {
self.put_aside()?;
self.unreadable = false;
}
if let Some(parent) = self.file.parent() {
std::fs::create_dir_all(parent)?;
std::fs::set_permissions(parent, std::fs::Permissions::from_mode(0o700))?;
}
let text = serde_json::to_string_pretty(&conversation.to_json()).map_err(io::Error::other)?;
atomic_write(&self.file, text.as_bytes())?;
std::fs::set_permissions(&self.file, std::fs::Permissions::from_mode(0o600))
}
pub fn put_aside(&self) -> io::Result<()> {
if !self.file.exists() {
return Ok(());
}
let seconds = SystemTime::now().duration_since(UNIX_EPOCH).map_or(0, |elapsed| elapsed.as_secs());
let stem = self.file.file_stem().map(|stem| stem.to_string_lossy().into_owned()).unwrap_or_default();
let mut aside = self.file.with_file_name(format!("{stem}-{seconds}.json"));
let mut extra = 1;
while aside.exists() {
aside = self.file.with_file_name(format!("{stem}-{seconds}-{extra}.json"));
extra += 1;
}
std::fs::rename(&self.file, aside)
}
}
fn fnv(bytes: &[u8]) -> u64 {
bytes.iter().fold(0xcbf2_9ce4_8422_2325, |hash, byte| (hash ^ u64::from(*byte)).wrapping_mul(0x0100_0000_01b3))
}
#[cfg(test)]
mod tests {
use super::*;
fn root(name: &str) -> PathBuf {
let root = std::env::temp_dir().join(format!("qcli-store-{name}-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&root);
root
}
fn store_in(root: &Path, folder: &str) -> Store {
std::fs::create_dir_all(root).expect("folder");
Store::new(root, Path::new(folder))
}
fn mode(path: &Path) -> u32 {
use std::os::unix::fs::PermissionsExt as _;
std::fs::metadata(path).expect("there").permissions().mode() & 0o777
}
#[test]
fn conversations_are_kept_where_only_their_owner_can_read_them() {
use std::os::unix::fs::PermissionsExt as _;
let base = root("private");
let conversations = base.join("conversations");
std::fs::create_dir_all(&conversations).expect("folder");
std::fs::set_permissions(&conversations, std::fs::Permissions::from_mode(0o755)).expect("open, as before");
let mut store = Store::new(&conversations, Path::new("/work/secret"));
let mut conversation = Conversation::default();
conversation.say("the token is in .env");
store.save(&conversation).expect("saved");
assert_eq!(mode(&conversations), 0o700, "the folder is its owner's alone");
assert_eq!(mode(store.file()), 0o600, "and so is the file");
let fresh = base.join("fresh").join("conversations");
let mut store = Store::new(&fresh, Path::new("/work/secret"));
store.save(&conversation).expect("saved");
assert_eq!(mode(&fresh), 0o700);
assert_eq!(mode(store.file()), 0o600);
let _ = std::fs::remove_dir_all(&base);
}
#[test]
fn a_folder_gets_its_conversation_back_and_two_folders_of_one_name_do_not_share() {
let root = root("round");
let mut store = store_in(&root, "/work/my project");
assert!(store.load().is_empty(), "nothing yet");
let mut conversation = Conversation::default();
conversation.say("hello");
store.save(&conversation).expect("saved");
let mut again = Store::new(&root, Path::new("/work/my project"));
assert_eq!(again.load(), conversation);
assert!(store.file().file_name().is_some_and(|name| name.to_string_lossy().starts_with("my_project-")));
assert_ne!(Store::new(&root, Path::new("/other/my project")).file(), store.file());
std::fs::remove_dir_all(&root).expect("clean up");
}
#[test]
fn a_new_conversation_puts_the_old_one_aside_and_a_broken_file_reads_as_empty() {
let root = root("aside");
let mut store = store_in(&root, "/work/a");
let mut conversation = Conversation::default();
conversation.say("keep me");
store.save(&conversation).expect("saved");
store.put_aside().expect("moved");
assert!(store.load().is_empty());
let kept: Vec<_> = std::fs::read_dir(&root).expect("folder").collect();
assert_eq!(kept.len(), 1, "the old conversation is still there");
std::fs::write(store.file(), "{ not json").expect("write");
assert!(store.load().is_empty());
std::fs::remove_dir_all(&root).expect("clean up");
}
const NOT_CONVERSATIONS: [(&str, &[u8]); 6] = [
("empty", b""),
("one-byte", b"x"),
("latin1", &[0xff, 0xfe, 0x7b, 0x7d]),
("list", b"[1,2,3]"),
("messages-number", br#"{"messages": 5}"#),
("half-json", br#"{"messages": [{"role": "user", "content": [{"type": "text", "text": "cut off here"#),
];
#[test]
fn a_file_that_is_not_a_conversation_reads_as_empty_and_is_never_written_over() {
for (name, bytes) in NOT_CONVERSATIONS {
let root = root(&format!("not-{name}"));
let mut store = store_in(&root, "/work/a");
std::fs::write(store.file(), bytes).expect("write");
assert!(store.load().is_empty(), "{name}: nothing to show");
let mut conversation = Conversation::default();
conversation.say("a new thread");
store.save(&conversation).expect("saved");
let others: Vec<PathBuf> = std::fs::read_dir(&root)
.expect("folder")
.map(|entry| entry.expect("an entry").path())
.filter(|path| path != store.file())
.collect();
assert_eq!(others.len(), 1, "{name}: the file is kept, not overwritten");
assert_eq!(std::fs::read(&others[0]).expect("read"), bytes, "{name}: the very same bytes");
let mut again = Store::new(&root, Path::new("/work/a"));
assert_eq!(again.load(), conversation, "{name}: the new thread is saved in its place");
std::fs::remove_dir_all(&root).expect("clean up");
}
}
#[test]
fn a_saved_conversation_is_written_over_plainly_the_next_time() {
let root = root("save-twice");
let mut store = store_in(&root, "/work/a");
std::fs::write(store.file(), "not a conversation").expect("write");
assert!(store.load().is_empty());
let mut conversation = Conversation::default();
conversation.say("first");
store.save(&conversation).expect("saved");
conversation.say("second");
store.save(&conversation).expect("saved again");
assert_eq!(std::fs::read_dir(&root).expect("folder").count(), 2, "the kept file and the conversation");
assert_eq!(store.load(), conversation, "the last save is the file now");
std::fs::remove_dir_all(&root).expect("clean up");
}
#[test]
fn a_hundred_thousand_messages_come_back_one_by_one() {
let root = root("many-messages");
let mut store = store_in(&root, "/work/a");
let mut file = std::fs::File::create(store.file()).expect("create");
{
use std::io::Write as _;
file.write_all(br#"{"version": 1, "messages": ["#).expect("write");
for number in 0..100_000 {
let message = format!(r#"{{"role": "user", "content": [{{"type": "text", "text": "m{number}"}}]}}"#);
file.write_all(message.as_bytes()).expect("write");
file.write_all(if number + 1 == 100_000 { b"]}" } else { b"," }).expect("write");
}
}
drop(file);
let conversation = store.load();
assert_eq!(conversation.messages().len(), 100_000, "every message of the file");
assert_eq!(conversation.messages()[0]["content"][0]["text"], "m0");
assert_eq!(conversation.messages()[99_999]["content"][0]["text"], "m99999");
std::fs::remove_dir_all(&root).expect("clean up");
}
#[test]
fn a_ten_megabyte_conversation_file_comes_back_whole() {
let root = root("ten-megabytes");
let mut store = store_in(&root, "/work/a");
let piece = "y".repeat(64 * 1024);
let mut file = std::fs::File::create(store.file()).expect("create");
{
use std::io::Write as _;
file.write_all(br#"{"version": 1, "messages": [{"role": "user", "content": [{"type": "text", "text": ""#)
.expect("write");
for _ in 0..160 {
file.write_all(piece.as_bytes()).expect("write");
}
file.write_all(b"\"}]}]}").expect("write");
}
drop(file);
assert!(std::fs::metadata(store.file()).expect("the file").len() >= 10 * 1024 * 1024, "ten megabytes");
let conversation = store.load();
let text = conversation.messages()[0]["content"][0]["text"].as_str().expect("the text");
assert_eq!(text.len(), 160 * piece.len(), "nothing of it was cut");
std::fs::remove_dir_all(&root).expect("clean up");
}
#[test]
fn one_message_of_five_thousand_calls_comes_back_with_every_call_answered() {
let root = root("many-blocks");
let mut store = store_in(&root, "/work/a");
let mut file = String::from(r#"{"version": 1, "messages": [{"role": "assistant", "content": ["#);
for number in 0..5000 {
file.push_str(&format!(
r#"{{"type": "tool_use", "id": "c{number}", "name": "list_dir", "input": {{"path": "."}}}}"#
));
if number + 1 < 5000 {
file.push(',');
}
}
file.push_str("]}]}");
std::fs::write(store.file(), file).expect("write");
let conversation = store.load();
assert!(conversation.open_calls().is_empty(), "no call is left waiting");
let answered = conversation.messages().last().expect("the results");
let ids: Vec<_> = crate::conversation::blocks(answered)
.map(|block| block["tool_use_id"].as_str().unwrap_or_default().to_owned())
.collect();
assert_eq!(ids.len(), 5000, "every call of the reply is answered");
assert_eq!((ids.first().map(String::as_str), ids.last().map(String::as_str)), (Some("c0"), Some("c4999")));
std::fs::remove_dir_all(&root).expect("clean up");
}
}