pub use kcode_k1_chat_persistence_records::{
Batch, BoxId, CODEC_VERSION, ChatBox, Error, EventRecord, Record, SessionId, SessionLog,
ToolCallId, TxId,
};
use kcode_k1_chat_persistence_records::{
encode_line, locate_predecessor, strict_prefix, to_log, validate_session_records,
};
use std::fs::{self, File, OpenOptions};
use std::io::{self, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
type Result<T, E = Error> = std::result::Result<T, E>;
pub struct Projection {
root: PathBuf,
sessions: PathBuf,
}
impl Projection {
pub fn new(root: impl AsRef<Path>) -> Result<Self> {
let root = root.as_ref().to_path_buf();
fs::create_dir_all(&root)?;
let sessions = root.join("sessions");
fs::create_dir_all(&sessions)?;
Ok(Self { root, sessions })
}
pub fn load(&self, session: SessionId) -> Result<SessionLog> {
let records = self
.read_strict(session)?
.into_iter()
.map(|value| value.1)
.collect::<Vec<_>>();
validate_session_records(session, &records)?;
Ok(to_log(records))
}
pub fn apply(&mut self, txid: TxId, batch: &Batch, reconcile_first: bool) -> Result<()> {
batch.encode()?;
fs::create_dir_all(&self.sessions)?;
let path = self.path(batch.session_id);
if reconcile_first {
let bytes = match fs::read(&path) {
Ok(value) => value,
Err(error) if error.kind() == io::ErrorKind::NotFound => Vec::new(),
Err(error) => return Err(error.into()),
};
let end = locate_predecessor(&bytes, batch.predecessor.as_ref())?;
let mut combined = strict_prefix(&bytes[..end])?
.into_iter()
.map(|value| value.1)
.collect::<Vec<_>>();
combined.extend(batch.records.clone());
validate_session_records(batch.session_id, &combined)?;
let mut file = OpenOptions::new()
.create(true)
.read(true)
.write(true)
.truncate(false)
.open(&path)?;
file.set_len(end as u64)?;
file.seek(SeekFrom::Start(end as u64))?;
append_lines(&mut file, txid, &batch.records)?;
file.sync_all()?;
} else {
let current = self.read_strict(batch.session_id)?;
let mut combined = current
.iter()
.map(|value| value.1.clone())
.collect::<Vec<_>>();
if combined.last() != batch.predecessor.as_ref() {
return Err(Error::Invalid("predecessor mismatch"));
}
combined.extend(batch.records.clone());
validate_session_records(batch.session_id, &combined)?;
let mut file = OpenOptions::new().create(true).append(true).open(&path)?;
append_lines(&mut file, txid, &batch.records)?;
file.sync_all()?;
}
Ok(())
}
pub fn discard_all(&mut self) -> Result<()> {
match fs::remove_dir_all(&self.sessions) {
Ok(()) => {}
Err(error) if error.kind() == io::ErrorKind::NotFound => {}
Err(error) => return Err(error.into()),
}
File::open(&self.root)?.sync_all()?;
Ok(())
}
fn path(&self, session: SessionId) -> PathBuf {
self.sessions.join(format!("{}.jsonl", hex(session)))
}
fn read_strict(&self, session: SessionId) -> Result<Vec<(TxId, Record, usize)>> {
let bytes = match fs::read(self.path(session)) {
Ok(value) => value,
Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(Vec::new()),
Err(error) => return Err(error.into()),
};
strict_prefix(&bytes)
}
}
fn append_lines(file: &mut File, txid: TxId, records: &[Record]) -> Result<()> {
for record in records {
file.write_all(&encode_line(txid, record)?)?;
}
Ok(())
}
fn hex(bytes: [u8; 12]) -> String {
bytes.iter().map(|value| format!("{value:02x}")).collect()
}
#[cfg(test)]
mod tests {
use super::*;
use kcode_k1_chat_chatend::{Chatend, ProviderCall, ResultView, ToolResult, ToolResultStatus};
use serde_json::json;
use std::sync::atomic::{AtomicU64, Ordering};
fn directory() -> PathBuf {
static NEXT: AtomicU64 = AtomicU64::new(0);
let path = std::env::temp_dir().join(format!(
"k1-persistence-store-{}-{}",
std::process::id(),
NEXT.fetch_add(1, Ordering::Relaxed)
));
let _ = fs::remove_dir_all(&path);
path
}
fn tx(value: u8) -> TxId {
TxId::from_bytes([value; 12])
}
fn current_records() -> Vec<Record> {
let mut chat = Chatend::new();
chat.accept_box(
"Future".into(),
"opaque".into(),
"future/v9".into(),
"hidden".into(),
)
.unwrap();
chat.start_round().unwrap();
let call = chat
.append_stage(
String::new(),
vec![ProviderCall {
tool: "Work".into(),
tool_version: "1.0.0".into(),
arguments: json!({"nested": [1, {"ok": true}]}),
}],
)
.unwrap()
.remove(0);
assert_eq!(call.call.call_id().to_string(), "c1");
chat.accept_async_return(
ToolResult::new(
call.call.call_id(),
call.call_box_id,
"Work".into(),
"1.0.0".into(),
ToolResultStatus::Ok,
json!({"done": true}),
ResultView::OneLine("success".into()),
)
.unwrap(),
)
.unwrap();
chat.done(String::new()).unwrap();
chat.boxes().iter().cloned().map(Record::Box).collect()
}
#[test]
fn records_facade_and_current_transcript_round_trip() {
fn assert_type<T>() {}
assert_type::<Batch>();
assert_type::<BoxId>();
assert_type::<ChatBox>();
assert_type::<EventRecord>();
assert_type::<Record>();
assert_type::<SessionId>();
assert_type::<SessionLog>();
assert_type::<ToolCallId>();
assert_type::<TxId>();
assert_eq!(CODEC_VERSION, 2);
let mut records = current_records();
records.push(Record::Event(
EventRecord::new(3, 1, 0, "handler".into(), json!({"saved": true})).unwrap(),
));
let batch = Batch::new([1; 12], None, records.clone()).unwrap();
let root = directory();
let mut projection = Projection::new(&root).unwrap();
projection.apply(tx(0xab), &batch, false).unwrap();
assert_eq!(projection.load([1; 12]).unwrap().records, records);
let bytes = fs::read(projection.path([1; 12])).unwrap();
assert!(
bytes
.windows(b"future/v9".len())
.any(|value| value == b"future/v9")
);
fs::remove_dir_all(root).unwrap();
}
#[test]
fn append_reconcile_restart_and_partial_suffix() {
let root = directory();
let mut projection = Projection::new(&root).unwrap();
let records = current_records();
let first = Batch::new([4; 12], None, records[..2].to_vec()).unwrap();
projection.apply(tx(1), &first, false).unwrap();
let next = Batch::new([4; 12], Some(records[1].clone()), records[2..].to_vec()).unwrap();
projection.apply(tx(2), &next, false).unwrap();
projection.apply(tx(2), &next, true).unwrap();
drop(projection);
let mut projection = Projection::new(&root).unwrap();
assert_eq!(projection.load([4; 12]).unwrap().records, records);
let mut file = OpenOptions::new()
.append(true)
.open(projection.path([4; 12]))
.unwrap();
file.write_all(b"{partial").unwrap();
projection.apply(tx(2), &next, true).unwrap();
assert_eq!(projection.load([4; 12]).unwrap().records.len(), 3);
projection.discard_all().unwrap();
assert!(projection.load([4; 12]).unwrap().records.is_empty());
fs::remove_dir_all(root).unwrap();
}
#[test]
fn malformed_lines_and_predecessors_fail_closed() {
let root = directory();
let mut projection = Projection::new(&root).unwrap();
fs::write(projection.path([1; 12]), b"{}\n").unwrap();
assert!(projection.load([1; 12]).is_err());
fs::write(projection.path([2; 12]), b"{}").unwrap();
assert!(projection.load([2; 12]).is_err());
let records = current_records();
let missing = Batch::new([3; 12], Some(records[1].clone()), records[2..].to_vec()).unwrap();
assert!(projection.apply(tx(3), &missing, true).is_err());
fs::remove_dir_all(root).unwrap();
}
}