use std::collections::{BTreeMap, HashMap, HashSet};
use std::io;
use std::ops::ControlFlow;
use rig_core::conversation::{ChangeRecord, Conversation};
use serde_json::Value;
use crate::blobs::{
BLOBS, load_images, load_images_with, messages_mut, named_blobs, put_blob, read_blob,
};
use crate::format::{Body, Header, LogLines, Meta, SESSION, parse, stamp_of};
use crate::store::{check_leaf, now};
use crate::writer::Book;
use crate::{FORMAT, LogRewrite, Repair, SessionError, Stamp, Store};
#[derive(Clone, Debug, Default, PartialEq)]
pub struct Session {
pub agents: Vec<SavedAgent>,
pub saved: BTreeMap<String, Value>,
pub repairs: Vec<Repair>,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct SavedAgent {
pub id: String,
pub parent: Option<String>,
pub started_ms: u64,
pub conversation: Conversation,
pub saved: BTreeMap<String, Value>,
pub retired: bool,
pub rewrites: Vec<LogRewrite>,
}
#[derive(Debug)]
pub struct Unreadable {
pub log: String,
pub error: SessionError,
}
#[derive(Clone, Debug, PartialEq)]
pub enum LogRecord {
Change(ChangeRecord),
Saved {
key: String,
value: Option<Value>,
},
Retired,
}
#[derive(Default)]
pub(crate) struct Folded {
pub name: String,
pub header: Header,
pub started_ms: u64,
pub conversation: Conversation,
pub state: BTreeMap<String, Value>,
pub retired: bool,
pub next_seq: u64,
pub next_n: u64,
pub len: u64,
pub torn: Option<(u64, u64)>,
pub lost: Vec<io::Error>,
}
struct Read {
header: Header,
started: Stamp,
last: Stamp,
}
pub(crate) fn log_lines(store: &dyn Store, name: &str) -> Result<Option<LogLines>, SessionError> {
match store.read_log(name) {
Err(failure) if failure.kind() == io::ErrorKind::NotFound => Ok(None),
bytes => Ok(Some(LogLines::parse(bytes?))),
}
}
pub(crate) fn read_logs(store: &dyn Store) -> Result<Vec<(String, LogLines)>, SessionError> {
let mut logs = Vec::new();
for name in store.logs()? {
logs.extend(log_lines(store, &name)?.map(|lines| (name, lines)));
}
Ok(logs)
}
fn read(
store: &dyn Store,
name: &str,
each: impl FnMut(Stamp, Body) -> Result<ControlFlow<()>, String>,
) -> Result<Option<Read>, SessionError> {
match log_lines(store, name)? {
Some(lines) => read_lines(name, &lines, each),
None => Ok(None),
}
}
fn read_lines(
name: &str,
lines: &LogLines,
mut each: impl FnMut(Stamp, Body) -> Result<ControlFlow<()>, String>,
) -> Result<Option<Read>, SessionError> {
let Some((_, _, first)) = lines.iter().next() else {
return Ok(None);
};
let refuse = |line, why| SessionError::Record {
log: name.to_owned(),
line,
why,
};
let (started, header) = read_header(name, first)?;
let mut last = started;
for (line, _, bytes) in lines.iter().skip(1) {
let (stamp, body) = parse(bytes).map_err(|why| refuse(line, why))?;
if matches!(body, Body::Meta(Meta::Header(_))) {
return Err(refuse(line, "a second header".to_owned()));
}
if each(stamp, body)
.map_err(|why| refuse(line, why))?
.is_break()
{
break;
}
last = stamp;
}
Ok(Some(Read {
header,
started,
last,
}))
}
pub(crate) fn read_header(name: &str, first: &[u8]) -> Result<(Stamp, Header), SessionError> {
let (started, header) = match parse(first) {
Ok((stamp, Body::Meta(Meta::Header(header)))) => (stamp, header),
_ => {
let found = serde_json::from_slice::<Header>(first)
.ok()
.and_then(|header| header.format);
return Err(SessionError::Format {
log: name.to_owned(),
found,
});
}
};
if header.format.as_deref() != Some(FORMAT) {
let found = header.format;
let log = name.to_owned();
return Err(SessionError::Format { log, found });
}
if !header.lineage.is_log() {
return Err(SessionError::Record {
log: name.to_owned(),
line: 1,
why: "a header without the log's lineage".to_owned(),
});
}
Ok((started, header))
}
#[derive(Clone, Copy, Default)]
pub(crate) struct Fold {
pub until: Option<u64>,
pub images: bool,
}
impl Fold {
pub const WHOLE: Self = Self {
until: None,
images: true,
};
}
pub(crate) fn fold(
store: &dyn Store,
name: &str,
how: Fold,
) -> Result<Option<Folded>, SessionError> {
match log_lines(store, name)? {
Some(lines) => fold_lines(name, &lines, how, &|blob| read_blob(store, blob)),
None => Ok(None),
}
}
pub(crate) fn fold_lines(
name: &str,
lines: &LogLines,
how: Fold,
blob: &dyn Fn(&str) -> io::Result<Vec<u8>>,
) -> Result<Option<Folded>, SessionError> {
let mut conversation = Conversation::new();
let mut state = BTreeMap::new();
let mut retired = false;
let mut lost = Vec::new();
let read = read_lines(name, lines, |stamp, body| {
if how.until.is_some_and(|until| stamp.n >= until) {
return Ok(ControlFlow::Break(()));
}
match body {
Body::Meta(meta) => note_meta(meta, &mut state, &mut retired),
Body::Change(mut recorded) => {
if how.images {
let _ = messages_mut(&mut recorded.change, &mut |message| {
if let Err(failure) = load_images_with(message, blob) {
lost.push(failure);
}
Ok(())
});
}
conversation
.apply(&recorded)
.map_err(|error| error.to_string())?;
}
}
Ok(ControlFlow::Continue(()))
})?;
let Some(read) = read else {
return Ok(None);
};
let changes = conversation.revision();
conversation.continue_log(read.header.lineage, changes);
let tail = lines.tail().len() as u64;
Ok(Some(Folded {
name: name.to_owned(),
started_ms: read.started.t,
header: read.header,
conversation,
state,
retired,
next_seq: read.last.seq.saturating_add(1),
next_n: read.last.n.saturating_add(1),
len: lines.len(),
torn: (tail > 0).then(|| (lines.len(), tail)),
lost,
}))
}
fn note_meta(meta: Meta, state: &mut BTreeMap<String, Value>, retired: &mut bool) {
match meta {
Meta::State { key, value } => {
match value.filter(|value| !value.is_null()) {
Some(value) => state.insert(key, value),
None => state.remove(&key),
};
}
Meta::Retired => *retired = true,
Meta::Header(_) => {}
}
}
fn tally(name: &str, lines: &LogLines) -> Result<Option<Folded>, SessionError> {
let Some((_, _, first)) = lines.iter().next() else {
return Ok(None);
};
let (started, header) = read_header(name, first)?;
let (mut last, mut changes) = (started, 0);
let (mut state, mut retired) = (BTreeMap::new(), false);
for (line, _, bytes) in lines.iter().skip(1) {
let refuse = |why| SessionError::Record {
log: name.to_owned(),
line,
why,
};
let (stamp, change) = stamp_of(bytes).map_err(refuse)?;
if !change && let (_, Body::Meta(meta)) = parse(bytes).map_err(refuse)? {
note_meta(meta, &mut state, &mut retired);
}
last = stamp;
changes += u64::from(change);
}
let mut conversation = Conversation::new();
conversation.continue_log(header.lineage, changes);
let tail = lines.tail().len() as u64;
Ok(Some(Folded {
name: name.to_owned(),
started_ms: started.t,
header,
conversation,
state,
retired,
next_seq: last.seq.saturating_add(1),
next_n: last.n.saturating_add(1),
len: lines.len(),
torn: (tail > 0).then(|| (lines.len(), tail)),
..Folded::default()
}))
}
pub(crate) fn fold_all(
store: &dyn Store,
logs: &[(String, LogLines)],
) -> Result<Vec<Folded>, SessionError> {
let mut folded = Vec::new();
for (name, lines) in logs {
folded.extend(fold_lines(name, lines, Fold::WHOLE, &|blob| {
read_blob(store, blob)
})?);
}
Ok(folded)
}
pub(crate) fn repairs(store: &dyn Store, folded: &[Folded]) -> Result<Vec<Repair>, SessionError> {
let torn = torn_tails(store, folded)?.into_iter();
let mut repairs: Vec<Repair> = torn
.map(|(log, kept, cut)| Repair::TornTail {
log,
kept,
cut,
set_aside: None,
})
.collect();
for log in folded {
repairs.extend(log.lost.iter().map(|why| Repair::LostImage {
log: log.name.clone(),
why: why.to_string(),
}));
}
Ok(repairs)
}
pub(crate) fn torn_tails(
store: &dyn Store,
folded: &[Folded],
) -> Result<Vec<(String, u64, u64)>, SessionError> {
let mut torn = Vec::new();
for log in store.logs()? {
match folded.iter().find(|folded| folded.name == log) {
Some(folded) => torn.extend(folded.torn.map(|(kept, cut)| (log, kept, cut))),
None => torn.extend(torn_whole(store, &log).map(|cut| (log, 0, cut))),
}
}
Ok(torn)
}
fn torn_whole(store: &dyn Store, log: &str) -> Option<u64> {
let lines = log_lines(store, log).ok().flatten();
let lines = lines.filter(|lines| lines.len() == 0)?;
let cut = lines.tail().len() as u64;
(cut > 0).then_some(cut)
}
pub(crate) fn retired(logs: &[Folded]) -> HashSet<String> {
let logs = logs.iter().map(|log| {
let parent = log.header.parent.as_deref();
(log.name.as_str(), parent, log.retired)
});
skipped(logs).into_iter().map(str::to_owned).collect()
}
pub(crate) fn skipped<'a>(
logs: impl Iterator<Item = (&'a str, Option<&'a str>, bool)>,
) -> HashSet<&'a str> {
let logs: Vec<_> = logs.collect();
let parents: HashMap<&str, Option<&str>> = logs
.iter()
.map(|&(name, parent, _)| (name, parent))
.collect();
let gone: HashSet<&str> = logs
.iter()
.filter(|&&(_, _, retired)| retired)
.map(|&(name, ..)| name)
.collect();
let mut skipped = HashSet::new();
for &(name, ..) in &logs {
let up = std::iter::successors(Some(name), |id| parents.get(id).copied().flatten());
if up.take(logs.len() + 1).any(|id| gone.contains(id)) {
skipped.insert(name);
}
}
skipped
}
pub(crate) fn session_of(folded: Vec<Folded>, repairs: Vec<Repair>) -> Session {
let mut session = Session {
repairs,
..Session::default()
};
for log in folded {
if log.name == SESSION {
session.saved = log.state;
continue;
}
session.agents.push(SavedAgent {
id: log.name,
parent: log.header.parent,
started_ms: log.started_ms,
conversation: log.conversation,
saved: log.state,
retired: log.retired,
rewrites: log.header.rewrites,
});
}
session
.agents
.sort_by(|a, b| (a.started_ms, &a.id).cmp(&(b.started_ms, &b.id)));
session
}
impl Session {
pub fn read(store: &dyn Store) -> Result<Self, SessionError> {
let folded = fold_all(store, &read_logs(store)?)?;
let repairs = repairs(store, &folded)?;
Ok(session_of(folded, repairs))
}
pub fn read_lenient(store: &dyn Store) -> Result<(Self, Vec<Unreadable>), SessionError> {
let mut folded = Vec::new();
let mut unreadable = Vec::new();
for log in store.logs()? {
match fold(store, &log, Fold::WHOLE) {
Ok(read) => folded.extend(read),
Err(error) => unreadable.push(Unreadable { log, error }),
}
}
let repairs = repairs(store, &folded)?;
Ok((session_of(folded, repairs), unreadable))
}
pub fn agent(&self, id: &str) -> Option<&SavedAgent> {
self.agents.iter().find(|agent| agent.id == id)
}
pub fn agent_mut(&mut self, id: &str) -> Option<&mut SavedAgent> {
self.agents.iter_mut().find(|agent| agent.id == id)
}
pub fn save(&mut self, store: &dyn Store) -> Result<(), SessionError> {
let mut folded = Vec::new();
for (name, lines) in &read_logs(store)? {
folded.extend(tally(name, lines)?);
}
let had = |id: &str| folded.iter().find(|log| log.name == id);
for agent in &self.agents {
if agent.id == SESSION {
let why = format!("{SESSION:?} is the session's own log, not an agent's");
return Err(io::Error::new(io::ErrorKind::InvalidInput, why).into());
}
check_continues(had(&agent.id), &agent.id, &agent.conversation)?;
}
let mut book = Book::after(&folded);
let torn = torn_tails(store, &folded)?.into_iter();
book.set_aside(torn.map(|(log, _, cut)| (log, cut)));
book.start();
for agent in &mut self.agents {
let _ = book.opened(&agent.id, agent.parent.as_deref());
let put = &mut |name: String, bytes: Vec<u8>| put_blob(store, &name, &bytes);
book.save(put, &agent.id, &mut agent.conversation);
let known = had(&agent.id).map(|log| &log.state);
book.states(&agent.id, known, &agent.saved);
if agent.retired && !had(&agent.id).is_some_and(|log| log.retired) {
book.retire(&agent.id);
}
}
let known = had(SESSION).map(|log| &log.state);
book.states(SESSION, known, &self.saved);
book.flush(store)
}
}
fn check_continues(
log: Option<&Folded>,
name: &str,
conversation: &Conversation,
) -> Result<(), SessionError> {
let Some(log) = log else {
return Ok(());
};
let changes = log.conversation.revision();
if changes > 0 && !conversation.continues(log.header.lineage, changes) {
return Err(SessionError::Changed {
log: name.to_owned(),
});
}
Ok(())
}
pub fn load(store: &dyn Store, log: &str) -> Result<Conversation, SessionError> {
check_leaf(log)?;
conversation_of(fold(store, log, Fold::WHOLE)?)
}
fn conversation_of(folded: Option<Folded>) -> Result<Conversation, SessionError> {
let Some(folded) = folded else {
return Ok(Conversation::new());
};
let mut lost = folded.lost.into_iter();
match lost.find(|failure| failure.kind() != io::ErrorKind::NotFound) {
Some(failure) => Err(failure.into()),
None => Ok(folded.conversation),
}
}
pub fn save(
store: &dyn Store,
log: &str,
conversation: &mut Conversation,
) -> Result<(), SessionError> {
now(save_async(store, log, conversation))
}
pub async fn load_async<S: crate::AsyncStore + ?Sized>(
store: &S,
log: &str,
) -> Result<Conversation, SessionError> {
check_leaf(log)?;
let bytes = match store.read_log_async(log).await {
Err(failure) if failure.kind() == io::ErrorKind::NotFound => {
return Ok(Conversation::new());
}
bytes => bytes?,
};
let mut named = std::collections::BTreeSet::new();
named_blobs(&bytes, &mut named);
let mut blobs = BTreeMap::new();
for name in named {
let read = store.read_file_async(&format!("{BLOBS}{name}")).await;
blobs.insert(name, read);
}
let blob = |name: &str| match blobs.get(name) {
Some(Ok(bytes)) => Ok(bytes.clone()),
Some(Err(failure)) => Err(io::Error::new(failure.kind(), failure.to_string())),
None => Err(io::ErrorKind::NotFound.into()),
};
conversation_of(fold_lines(
log,
&LogLines::parse(bytes),
Fold::WHOLE,
&blob,
)?)
}
pub async fn save_async<S: crate::AsyncStore + ?Sized>(
store: &S,
log: &str,
conversation: &mut Conversation,
) -> Result<(), SessionError> {
let lines = match store.read_log_async(log).await {
Err(failure) if failure.kind() == io::ErrorKind::NotFound => LogLines::parse(Vec::new()),
bytes => LogLines::parse(bytes?),
};
let folded = tally(log, &lines)?;
check_continues(folded.as_ref(), log, conversation)?;
let mut book = Book::after(folded.as_slice());
let cut = lines.tail().len() as u64;
book.set_aside((cut > 0).then(|| (log.to_owned(), cut)));
let mut staged = conversation.clone();
book.start();
let _ = book.opened(log, None);
let mut blobs = Vec::new();
let put = &mut |name, bytes| {
blobs.push((name, bytes));
Ok(())
};
book.save(put, log, &mut staged);
for (name, bytes) in blobs {
store
.write_file_async(&format!("{BLOBS}{name}"), &bytes)
.await?;
}
book.flush_async(store).await?;
*conversation = staged;
Ok(())
}
pub fn history(store: &dyn Store, log: &str) -> Result<Vec<(Stamp, LogRecord)>, SessionError> {
let mut records = Vec::new();
read(store, log, |stamp, body| {
let record = match body {
Body::Change(mut recorded) => {
let _ = messages_mut(&mut recorded.change, &mut |message| {
let _ = load_images(message, store);
Ok(())
});
LogRecord::Change(recorded)
}
Body::Meta(Meta::State { key, value }) => LogRecord::Saved {
key,
value: value.filter(|value| !value.is_null()),
},
Body::Meta(Meta::Retired) => LogRecord::Retired,
Body::Meta(Meta::Header(_)) => return Ok(ControlFlow::Continue(())),
};
records.push((stamp, record));
Ok(ControlFlow::Continue(()))
})?;
Ok(records)
}
pub fn as_of(store: &dyn Store, log: &str, n: u64) -> Result<Conversation, SessionError> {
let how = Fold {
until: Some(n),
..Fold::WHOLE
};
conversation_of(fold(store, log, how)?)
}
#[cfg(test)]
mod tests;