use kcode_k1_chat_chatend::{BoxContent, ChatBox, Chatend};
pub use kcode_k1_chat_chatend::{BoxId, ToolCallId};
pub use kcode_k1_transaction_id::TxId;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::fs::{self, File, OpenOptions};
use std::io::{self, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
pub type SessionId = [u8; 12];
pub const CODEC_VERSION: u32 = 1;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct EventRecord {
pub after_box_id: u64,
pub event_index: u64,
pub connected_box_id: u64,
pub handler: String,
pub data: Value,
}
impl EventRecord {
pub fn new(
after_box_id: u64,
event_index: u64,
connected_box_id: u64,
handler: String,
data: Value,
) -> Result<Self, Error> {
let value = Self {
after_box_id,
event_index,
connected_box_id,
handler,
data,
};
validate_event(&value)?;
Ok(value)
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Record {
Box(ChatBox),
Event(EventRecord),
}
impl Record {
pub fn chat_box(value: ChatBox) -> Self {
Self::Box(value)
}
pub fn event(value: EventRecord) -> Self {
Self::Event(value)
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Batch {
pub version: u32,
pub session_id: SessionId,
pub predecessor: Option<Record>,
pub records: Vec<Record>,
}
impl Batch {
pub fn new(
session_id: SessionId,
predecessor: Option<Record>,
records: Vec<Record>,
) -> Result<Self, Error> {
let value = Self {
version: CODEC_VERSION,
session_id,
predecessor,
records,
};
validate_batch(&value)?;
Ok(value)
}
pub fn encode(&self) -> Result<Vec<u8>, Error> {
validate_batch(self)?;
Ok(serde_json::to_vec(&WireBatch::from(self))?)
}
pub fn decode(bytes: &[u8]) -> Result<Self, Error> {
let wire: WireBatch = serde_json::from_slice(bytes)?;
let value = Self::try_from(wire)?;
validate_batch(&value)?;
Ok(value)
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct SessionLog {
pub boxes: Vec<ChatBox>,
pub events: Vec<EventRecord>,
pub records: Vec<Record>,
}
#[derive(Debug)]
pub enum Error {
Io(io::Error),
Json(serde_json::Error),
Invalid(&'static str),
}
impl std::fmt::Display for Error {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Io(e) => write!(f, "I/O error: {e}"),
Self::Json(e) => write!(f, "JSON error: {e}"),
Self::Invalid(e) => f.write_str(e),
}
}
}
impl std::error::Error for Error {}
impl From<io::Error> for Error {
fn from(value: io::Error) -> Self {
Self::Io(value)
}
}
impl From<serde_json::Error> for Error {
fn from(value: serde_json::Error) -> Self {
Self::Json(value)
}
}
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)?;
let records: Vec<Record> = records.into_iter().map(|v| v.1).collect();
validate_records(&records)?;
Ok(to_log(records))
}
pub fn apply(&mut self, txid: TxId, batch: &Batch, reconcile_first: bool) -> Result<()> {
validate_batch(batch)?;
fs::create_dir_all(&self.sessions)?;
let path = self.path(batch.session_id);
if reconcile_first {
let bytes = match fs::read(&path) {
Ok(v) => v,
Err(e) if e.kind() == io::ErrorKind::NotFound => Vec::new(),
Err(e) => return Err(e.into()),
};
let end = locate_predecessor(&bytes, batch.predecessor.as_ref())?;
let prefix = strict_prefix(&bytes[..end])?;
let mut combined: Vec<Record> = prefix.into_iter().map(|v| v.1).collect();
combined.extend(batch.records.clone());
validate_records(&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 records: Vec<Record> = current.iter().map(|v| v.1.clone()).collect();
if records.last() != batch.predecessor.as_ref() {
return Err(Error::Invalid("predecessor mismatch"));
}
let mut combined = records;
combined.extend(batch.records.clone());
validate_records(&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(e) if e.kind() == io::ErrorKind::NotFound => {}
Err(e) => return Err(e.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(v) => v,
Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(Vec::new()),
Err(e) => return Err(e.into()),
};
strict_prefix(&bytes)
}
}
fn append_lines(file: &mut File, txid: TxId, records: &[Record]) -> Result<()> {
for record in records {
let line = WireLine {
txid: hex(*txid.as_bytes()),
record: WireRecord::from(record),
};
let mut bytes = serde_json::to_vec(&line)?;
bytes.push(b'\n');
file.write_all(&bytes)?;
}
Ok(())
}
fn strict_prefix(bytes: &[u8]) -> Result<Vec<(TxId, Record, usize)>> {
if !bytes.is_empty() && bytes.last() != Some(&b'\n') {
return Err(Error::Invalid("line is not newline terminated"));
}
let mut out = Vec::new();
let mut start = 0;
while start < bytes.len() {
let relative = bytes[start..]
.iter()
.position(|b| *b == b'\n')
.ok_or(Error::Invalid("line is not newline terminated"))?;
let end = start + relative + 1;
let wire: WireLine = serde_json::from_slice(&bytes[start..end - 1])?;
let tx = TxId::from_bytes(parse_hex(&wire.txid)?);
out.push((tx, Record::try_from(wire.record)?, end));
start = end;
}
let records: Vec<Record> = out.iter().map(|v| v.1.clone()).collect();
validate_records(&records)?;
Ok(out)
}
fn locate_predecessor(bytes: &[u8], wanted: Option<&Record>) -> Result<usize> {
let Some(wanted) = wanted else { return Ok(0) };
let mut start = 0;
let mut found = None;
let mut prefix = Vec::new();
while start < bytes.len() {
let Some(relative) = bytes[start..].iter().position(|b| *b == b'\n') else {
if found.is_some() {
break;
} else {
return Err(Error::Invalid("predecessor absent before incomplete line"));
}
};
let end = start + relative + 1;
let parsed = serde_json::from_slice::<WireLine>(&bytes[start..end - 1])
.map_err(Error::from)
.and_then(|v| Record::try_from(v.record));
let record = match parsed {
Ok(v) => v,
Err(_) if found.is_some() => break,
Err(e) => return Err(e),
};
if same_identity(&record, wanted) {
if &record != wanted {
return Err(Error::Invalid("predecessor identity mismatch"));
}
if found.is_some() {
return Err(Error::Invalid("duplicate predecessor identity"));
}
found = Some(end);
}
if found.is_none() {
prefix.push(record);
validate_records(&prefix)?;
}
start = end;
}
found.ok_or(Error::Invalid("predecessor absent"))
}
fn same_identity(a: &Record, b: &Record) -> bool {
match (a, b) {
(Record::Box(a), Record::Box(b)) => a.id() == b.id(),
(Record::Event(a), Record::Event(b)) => {
(a.after_box_id, a.event_index, a.connected_box_id)
== (b.after_box_id, b.event_index, b.connected_box_id)
}
_ => false,
}
}
fn validate_batch(batch: &Batch) -> Result<()> {
if batch.version != CODEC_VERSION {
return Err(Error::Invalid("unsupported version"));
}
if batch.records.is_empty() {
return Err(Error::Invalid("empty batch"));
}
validate_suffix(batch.predecessor.as_ref(), &batch.records)?;
if batch.predecessor.is_none() {
validate_records(&batch.records)?;
}
Ok(())
}
fn validate_suffix(predecessor: Option<&Record>, suffix: &[Record]) -> Result<()> {
let (mut latest, mut event) = match predecessor {
None => (0, 0),
Some(Record::Box(v)) => (v.id().get(), 0),
Some(Record::Event(v)) => {
validate_event(v)?;
(v.after_box_id, v.event_index)
}
};
for record in suffix {
match record {
Record::Box(v) => {
if v.id().get()
!= latest
.checked_add(1)
.ok_or(Error::Invalid("box ID overflow"))?
{
return Err(Error::Invalid("noncontiguous box ID"));
}
latest = v.id().get();
event = 0;
}
Record::Event(v) => {
validate_event(v)?;
if v.after_box_id != latest
|| v.event_index
!= event
.checked_add(1)
.ok_or(Error::Invalid("event index overflow"))?
|| v.connected_box_id > latest
{
return Err(Error::Invalid("invalid event order or association"));
}
event = v.event_index;
}
}
}
Ok(())
}
fn validate_event(value: &EventRecord) -> Result<()> {
if value.handler.is_empty() {
Err(Error::Invalid("empty event handler"))
} else {
Ok(())
}
}
fn validate_records(records: &[Record]) -> Result<()> {
validate_suffix(None, records)?;
let boxes: Vec<ChatBox> = records
.iter()
.filter_map(|v| match v {
Record::Box(v) => Some(v.clone()),
_ => None,
})
.collect();
Chatend::recover(boxes).map_err(|_| Error::Invalid("invalid Chatend recovery"))?;
Ok(())
}
fn to_log(records: Vec<Record>) -> SessionLog {
let boxes = records
.iter()
.filter_map(|v| match v {
Record::Box(v) => Some(v.clone()),
_ => None,
})
.collect();
let events = records
.iter()
.filter_map(|v| match v {
Record::Event(v) => Some(v.clone()),
_ => None,
})
.collect();
SessionLog {
boxes,
events,
records,
}
}
fn hex(bytes: [u8; 12]) -> String {
bytes.iter().map(|v| format!("{v:02x}")).collect()
}
fn parse_hex(value: &str) -> Result<[u8; 12]> {
if value.len() != 24
|| !value
.bytes()
.all(|v| v.is_ascii_digit() || (b'a'..=b'f').contains(&v))
{
return Err(Error::Invalid("invalid lowercase 24-hex value"));
}
let mut out = [0; 12];
for (i, slot) in out.iter_mut().enumerate() {
*slot = u8::from_str_radix(&value[i * 2..i * 2 + 2], 16)
.map_err(|_| Error::Invalid("invalid hex"))?;
}
Ok(out)
}
#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct WireBatch {
version: u32,
session_id: String,
predecessor: Option<WireRecord>,
records: Vec<WireRecord>,
}
impl From<&Batch> for WireBatch {
fn from(v: &Batch) -> Self {
Self {
version: v.version,
session_id: hex(v.session_id),
predecessor: v.predecessor.as_ref().map(WireRecord::from),
records: v.records.iter().map(WireRecord::from).collect(),
}
}
}
impl TryFrom<WireBatch> for Batch {
type Error = Error;
fn try_from(v: WireBatch) -> Result<Self> {
Ok(Self {
version: v.version,
session_id: parse_hex(&v.session_id)?,
predecessor: v.predecessor.map(Record::try_from).transpose()?,
records: v
.records
.into_iter()
.map(Record::try_from)
.collect::<Result<_>>()?,
})
}
}
#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct WireLine {
txid: String,
record: WireRecord,
}
#[derive(Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
enum WireRecord {
Box {
id: u64,
content: WireContent,
},
Event {
after_box_id: u64,
event_index: u64,
connected_box_id: u64,
handler: String,
data: Value,
},
}
impl From<&Record> for WireRecord {
fn from(v: &Record) -> Self {
match v {
Record::Box(v) => Self::Box {
id: v.id().get(),
content: WireContent::from(v.content()),
},
Record::Event(v) => Self::Event {
after_box_id: v.after_box_id,
event_index: v.event_index,
connected_box_id: v.connected_box_id,
handler: v.handler.clone(),
data: v.data.clone(),
},
}
}
}
impl TryFrom<WireRecord> for Record {
type Error = Error;
fn try_from(v: WireRecord) -> Result<Self> {
Ok(match v {
WireRecord::Box { id, content } => {
Self::Box(ChatBox::new(BoxId::new(id), BoxContent::try_from(content)?))
}
WireRecord::Event {
after_box_id,
event_index,
connected_box_id,
handler,
data,
} => Self::Event(EventRecord::new(
after_box_id,
event_index,
connected_box_id,
handler,
data,
)?),
})
}
}
#[derive(Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
enum WireContent {
System {
text: String,
},
User {
text: String,
},
Kennedy {
text: String,
},
Attachment,
KtoolCall {
session: String,
sequence: u64,
name: String,
arguments: String,
},
KtoolReturn {
session: String,
sequence: u64,
originating_call: u64,
result: WireResult,
},
}
impl From<&BoxContent> for WireContent {
fn from(v: &BoxContent) -> Self {
match v {
BoxContent::System(text) => Self::System { text: text.clone() },
BoxContent::User(text) => Self::User { text: text.clone() },
BoxContent::Kennedy { text } => Self::Kennedy { text: text.clone() },
BoxContent::Attachment => Self::Attachment,
BoxContent::KtoolCall {
tool_call_id,
name,
arguments,
} => Self::KtoolCall {
session: hex(tool_call_id.session()),
sequence: tool_call_id.sequence(),
name: name.clone(),
arguments: arguments.clone(),
},
BoxContent::KtoolReturn {
tool_call_id,
originating_call,
result,
} => Self::KtoolReturn {
session: hex(tool_call_id.session()),
sequence: tool_call_id.sequence(),
originating_call: originating_call.get(),
result: WireResult::from(result),
},
}
}
}
impl TryFrom<WireContent> for BoxContent {
type Error = Error;
fn try_from(v: WireContent) -> Result<Self> {
Ok(match v {
WireContent::System { text } => Self::System(text),
WireContent::User { text } => Self::User(text),
WireContent::Kennedy { text } => Self::Kennedy { text },
WireContent::Attachment => Self::Attachment,
WireContent::KtoolCall {
session,
sequence,
name,
arguments,
} => Self::KtoolCall {
tool_call_id: ToolCallId::new(parse_hex(&session)?, sequence),
name,
arguments,
},
WireContent::KtoolReturn {
session,
sequence,
originating_call,
result,
} => Self::KtoolReturn {
tool_call_id: ToolCallId::new(parse_hex(&session)?, sequence),
originating_call: BoxId::new(originating_call),
result: result.into(),
},
})
}
}
#[derive(Serialize, Deserialize)]
#[serde(tag = "status", rename_all = "snake_case", deny_unknown_fields)]
enum WireResult {
Ok { value: String },
Err { value: String },
}
impl From<&std::result::Result<String, String>> for WireResult {
fn from(v: &std::result::Result<String, String>) -> Self {
match v {
Ok(value) => Self::Ok {
value: value.clone(),
},
Err(value) => Self::Err {
value: value.clone(),
},
}
}
}
impl From<WireResult> for std::result::Result<String, String> {
fn from(v: WireResult) -> Self {
match v {
WireResult::Ok { value } => Ok(value),
WireResult::Err { value } => Err(value),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use std::sync::atomic::{AtomicU64, Ordering};
fn bx(id: u64, content: BoxContent) -> Record {
Record::Box(ChatBox::new(BoxId::new(id), content))
}
fn event(after: u64, index: u64) -> Record {
Record::Event(
EventRecord::new(
after,
index,
after,
"handler".into(),
json!({"z":[true,null,{"a":1}]}),
)
.unwrap(),
)
}
fn dir() -> PathBuf {
static N: AtomicU64 = AtomicU64::new(0);
let p = std::env::temp_dir().join(format!(
"k1-persistence-{}-{}",
std::process::id(),
N.fetch_add(1, Ordering::Relaxed)
));
let _ = fs::remove_dir_all(&p);
p
}
fn tx(n: u8) -> TxId {
TxId::from_bytes([n; 12])
}
#[test]
fn codec_roundtrips_all_kinds_events_and_data() {
let tool = ToolCallId::new([9; 12], 7);
let records = vec![
event(0, 1),
bx(1, BoxContent::System("s".into())),
bx(2, BoxContent::User("u".into())),
bx(3, BoxContent::Kennedy { text: "k".into() }),
bx(4, BoxContent::Attachment),
bx(
5,
BoxContent::KtoolCall {
tool_call_id: tool,
name: "n".into(),
arguments: "{}".into(),
},
),
bx(
6,
BoxContent::KtoolReturn {
tool_call_id: tool,
originating_call: BoxId::new(5),
result: Err("e".into()),
},
),
event(6, 1),
];
let batch = Batch::new([1; 12], None, records).unwrap();
assert_eq!(Batch::decode(&batch.encode().unwrap()).unwrap(), batch);
assert!(Batch::decode(br#"{"version":1,"session_id":"010101010101010101010101","predecessor":null,"records":[],"extra":0}"#).is_err());
}
#[test]
fn order_and_chatend_validation_fail_closed() {
assert!(Batch::new([0; 12], None, vec![bx(2, BoxContent::Attachment)]).is_err());
assert!(EventRecord::new(0, 1, 0, String::new(), json!(null)).is_err());
assert!(
Batch::new(
[0; 12],
None,
vec![Record::Event(EventRecord {
after_box_id: 0,
event_index: 1,
connected_box_id: 1,
handler: "h".into(),
data: json!(null)
})]
)
.is_err()
);
let tool = ToolCallId::new([2; 12], 1);
assert!(
Batch::new(
[0; 12],
None,
vec![bx(
1,
BoxContent::KtoolReturn {
tool_call_id: tool,
originating_call: BoxId::new(1),
result: Ok("x".into())
}
)]
)
.is_err()
);
let valid = Batch::new(
[0; 12],
None,
vec![
bx(1, BoxContent::Attachment),
event(1, 1),
event(1, 2),
bx(2, BoxContent::User("next".into())),
],
)
.unwrap();
assert_eq!(valid.records.len(), 4);
}
#[test]
fn append_load_isolation_predecessors_and_strict_lines() {
let root = dir();
let mut p = Projection::new(&root).unwrap();
let one = Batch::new([1; 12], None, vec![bx(1, BoxContent::System("a".into()))]).unwrap();
p.apply(tx(1), &one, false).unwrap();
let pred = one.records.last().cloned();
let two = Batch::new(
[1; 12],
pred,
vec![event(1, 1), bx(2, BoxContent::User("b".into()))],
)
.unwrap();
p.apply(tx(2), &two, false).unwrap();
assert_eq!(p.load([1; 12]).unwrap().records.len(), 3);
assert!(p.load([2; 12]).unwrap().records.is_empty());
assert!(p.apply(tx(3), &one, false).is_err());
fs::write(p.path([2; 12]), b"{}\n").unwrap();
assert!(p.load([2; 12]).is_err());
fs::write(p.path([3; 12]), b"{}").unwrap();
assert!(p.load([3; 12]).is_err());
fs::remove_dir_all(root).unwrap();
}
#[test]
fn reconciliation_replaces_partial_and_fully_appended_callback() {
let root = dir();
let mut p = Projection::new(&root).unwrap();
let first = Batch::new([4; 12], None, vec![bx(1, BoxContent::System("a".into()))]).unwrap();
p.apply(tx(1), &first, false).unwrap();
let next = Batch::new(
[4; 12],
Some(first.records[0].clone()),
vec![event(1, 1), bx(2, BoxContent::User("b".into()))],
)
.unwrap();
p.apply(tx(2), &next, false).unwrap();
p.apply(tx(2), &next, true).unwrap();
assert_eq!(p.load([4; 12]).unwrap().records.len(), 3);
let path = p.path([4; 12]);
let mut f = OpenOptions::new().append(true).open(&path).unwrap();
f.write_all(b"{partial").unwrap();
p.apply(tx(2), &next, true).unwrap();
assert_eq!(p.load([4; 12]).unwrap().records.len(), 3);
let missing = Batch::new(
[4; 12],
Some(bx(9, BoxContent::Attachment)),
vec![bx(10, BoxContent::Attachment)],
)
.unwrap();
assert!(p.apply(tx(3), &missing, true).is_err());
fs::remove_dir_all(root).unwrap();
}
#[test]
fn discard_removes_only_sessions() {
let root = dir();
let keep = root.join("keep");
let mut p = Projection::new(&root).unwrap();
fs::write(&keep, b"x").unwrap();
let batch = Batch::new([5; 12], None, vec![bx(1, BoxContent::Attachment)]).unwrap();
p.apply(tx(1), &batch, false).unwrap();
p.discard_all().unwrap();
assert!(keep.exists() && !p.sessions.exists());
assert!(p.load([5; 12]).unwrap().records.is_empty());
fs::remove_dir_all(root).unwrap();
}
}