1use std::collections::{BTreeMap, HashMap, HashSet};
6use std::io;
7use std::ops::ControlFlow;
8
9use rig_core::conversation::{ChangeRecord, Conversation};
10use serde_json::Value;
11
12use crate::blobs::{
13 BLOBS, load_images, load_images_with, messages_mut, named_blobs, put_blob, read_blob,
14};
15use crate::format::{Body, Header, LogLines, Meta, SESSION, parse, stamp_of};
16use crate::store::{check_leaf, now};
17use crate::writer::Book;
18use crate::{FORMAT, LogRewrite, Repair, SessionError, Stamp, Store};
19
20#[derive(Clone, Debug, Default, PartialEq)]
22pub struct Session {
23 pub agents: Vec<SavedAgent>,
25 pub saved: BTreeMap<String, Value>,
27 pub repairs: Vec<Repair>,
29}
30
31#[derive(Clone, Debug, Default, PartialEq)]
33pub struct SavedAgent {
34 pub id: String,
36 pub parent: Option<String>,
38 pub started_ms: u64,
40 pub conversation: Conversation,
43 pub saved: BTreeMap<String, Value>,
45 pub retired: bool,
48 pub rewrites: Vec<LogRewrite>,
50}
51
52#[derive(Debug)]
54pub struct Unreadable {
55 pub log: String,
57 pub error: SessionError,
59}
60
61#[derive(Clone, Debug, PartialEq)]
63pub enum LogRecord {
64 Change(ChangeRecord),
66 Saved {
68 key: String,
70 value: Option<Value>,
72 },
73 Retired,
75}
76
77#[derive(Default)]
79pub(crate) struct Folded {
80 pub name: String,
81 pub header: Header,
82 pub started_ms: u64,
83 pub conversation: Conversation,
84 pub state: BTreeMap<String, Value>,
85 pub retired: bool,
86 pub next_seq: u64,
88 pub next_n: u64,
90 pub len: u64,
92 pub torn: Option<(u64, u64)>,
94 pub lost: Vec<io::Error>,
96}
97
98struct Read {
100 header: Header,
101 started: Stamp,
102 last: Stamp,
103}
104
105pub(crate) fn log_lines(store: &dyn Store, name: &str) -> Result<Option<LogLines>, SessionError> {
107 match store.read_log(name) {
108 Err(failure) if failure.kind() == io::ErrorKind::NotFound => Ok(None),
109 bytes => Ok(Some(LogLines::parse(bytes?))),
110 }
111}
112
113pub(crate) fn read_logs(store: &dyn Store) -> Result<Vec<(String, LogLines)>, SessionError> {
115 let mut logs = Vec::new();
116 for name in store.logs()? {
117 logs.extend(log_lines(store, &name)?.map(|lines| (name, lines)));
118 }
119 Ok(logs)
120}
121
122fn read(
124 store: &dyn Store,
125 name: &str,
126 each: impl FnMut(Stamp, Body) -> Result<ControlFlow<()>, String>,
127) -> Result<Option<Read>, SessionError> {
128 match log_lines(store, name)? {
129 Some(lines) => read_lines(name, &lines, each),
130 None => Ok(None),
131 }
132}
133
134fn read_lines(
138 name: &str,
139 lines: &LogLines,
140 mut each: impl FnMut(Stamp, Body) -> Result<ControlFlow<()>, String>,
141) -> Result<Option<Read>, SessionError> {
142 let Some((_, _, first)) = lines.iter().next() else {
143 return Ok(None);
144 };
145 let refuse = |line, why| SessionError::Record {
146 log: name.to_owned(),
147 line,
148 why,
149 };
150 let (started, header) = read_header(name, first)?;
151 let mut last = started;
152 for (line, _, bytes) in lines.iter().skip(1) {
153 let (stamp, body) = parse(bytes).map_err(|why| refuse(line, why))?;
154 if matches!(body, Body::Meta(Meta::Header(_))) {
155 return Err(refuse(line, "a second header".to_owned()));
156 }
157 if each(stamp, body)
158 .map_err(|why| refuse(line, why))?
159 .is_break()
160 {
161 break;
162 }
163 last = stamp;
164 }
165 Ok(Some(Read {
166 header,
167 started,
168 last,
169 }))
170}
171
172pub(crate) fn read_header(name: &str, first: &[u8]) -> Result<(Stamp, Header), SessionError> {
175 let (started, header) = match parse(first) {
176 Ok((stamp, Body::Meta(Meta::Header(header)))) => (stamp, header),
177 _ => {
178 let found = serde_json::from_slice::<Header>(first)
179 .ok()
180 .and_then(|header| header.format);
181 return Err(SessionError::Format {
182 log: name.to_owned(),
183 found,
184 });
185 }
186 };
187 if header.format.as_deref() != Some(FORMAT) {
188 let found = header.format;
189 let log = name.to_owned();
190 return Err(SessionError::Format { log, found });
191 }
192 if !header.lineage.is_log() {
193 return Err(SessionError::Record {
194 log: name.to_owned(),
195 line: 1,
196 why: "a header without the log's lineage".to_owned(),
197 });
198 }
199 Ok((started, header))
200}
201
202#[derive(Clone, Copy, Default)]
204pub(crate) struct Fold {
205 pub until: Option<u64>,
207 pub images: bool,
209}
210
211impl Fold {
212 pub const WHOLE: Self = Self {
214 until: None,
215 images: true,
216 };
217}
218
219pub(crate) fn fold(
223 store: &dyn Store,
224 name: &str,
225 how: Fold,
226) -> Result<Option<Folded>, SessionError> {
227 match log_lines(store, name)? {
228 Some(lines) => fold_lines(name, &lines, how, &|blob| read_blob(store, blob)),
229 None => Ok(None),
230 }
231}
232
233pub(crate) fn fold_lines(
236 name: &str,
237 lines: &LogLines,
238 how: Fold,
239 blob: &dyn Fn(&str) -> io::Result<Vec<u8>>,
240) -> Result<Option<Folded>, SessionError> {
241 let mut conversation = Conversation::new();
242 let mut state = BTreeMap::new();
243 let mut retired = false;
244 let mut lost = Vec::new();
245 let read = read_lines(name, lines, |stamp, body| {
246 if how.until.is_some_and(|until| stamp.n >= until) {
247 return Ok(ControlFlow::Break(()));
248 }
249 match body {
250 Body::Meta(meta) => note_meta(meta, &mut state, &mut retired),
251 Body::Change(mut recorded) => {
252 if how.images {
253 let _ = messages_mut(&mut recorded.change, &mut |message| {
254 if let Err(failure) = load_images_with(message, blob) {
255 lost.push(failure);
256 }
257 Ok(())
258 });
259 }
260 conversation
261 .apply(&recorded)
262 .map_err(|error| error.to_string())?;
263 }
264 }
265 Ok(ControlFlow::Continue(()))
266 })?;
267 let Some(read) = read else {
268 return Ok(None);
269 };
270 let changes = conversation.revision();
271 conversation.continue_log(read.header.lineage, changes);
272 let tail = lines.tail().len() as u64;
273 Ok(Some(Folded {
274 name: name.to_owned(),
275 started_ms: read.started.t,
276 header: read.header,
277 conversation,
278 state,
279 retired,
280 next_seq: read.last.seq.saturating_add(1),
281 next_n: read.last.n.saturating_add(1),
282 len: lines.len(),
283 torn: (tail > 0).then(|| (lines.len(), tail)),
284 lost,
285 }))
286}
287
288fn note_meta(meta: Meta, state: &mut BTreeMap<String, Value>, retired: &mut bool) {
291 match meta {
292 Meta::State { key, value } => {
293 match value.filter(|value| !value.is_null()) {
294 Some(value) => state.insert(key, value),
295 None => state.remove(&key),
296 };
297 }
298 Meta::Retired => *retired = true,
299 Meta::Header(_) => {}
300 }
301}
302
303fn tally(name: &str, lines: &LogLines) -> Result<Option<Folded>, SessionError> {
308 let Some((_, _, first)) = lines.iter().next() else {
309 return Ok(None);
310 };
311 let (started, header) = read_header(name, first)?;
312 let (mut last, mut changes) = (started, 0);
313 let (mut state, mut retired) = (BTreeMap::new(), false);
314 for (line, _, bytes) in lines.iter().skip(1) {
315 let refuse = |why| SessionError::Record {
316 log: name.to_owned(),
317 line,
318 why,
319 };
320 let (stamp, change) = stamp_of(bytes).map_err(refuse)?;
321 if !change && let (_, Body::Meta(meta)) = parse(bytes).map_err(refuse)? {
322 note_meta(meta, &mut state, &mut retired);
323 }
324 last = stamp;
325 changes += u64::from(change);
326 }
327 let mut conversation = Conversation::new();
328 conversation.continue_log(header.lineage, changes);
329 let tail = lines.tail().len() as u64;
330 Ok(Some(Folded {
331 name: name.to_owned(),
332 started_ms: started.t,
333 header,
334 conversation,
335 state,
336 retired,
337 next_seq: last.seq.saturating_add(1),
338 next_n: last.n.saturating_add(1),
339 len: lines.len(),
340 torn: (tail > 0).then(|| (lines.len(), tail)),
341 ..Folded::default()
342 }))
343}
344
345pub(crate) fn fold_all(
348 store: &dyn Store,
349 logs: &[(String, LogLines)],
350) -> Result<Vec<Folded>, SessionError> {
351 let mut folded = Vec::new();
352 for (name, lines) in logs {
353 folded.extend(fold_lines(name, lines, Fold::WHOLE, &|blob| {
354 read_blob(store, blob)
355 })?);
356 }
357 Ok(folded)
358}
359
360pub(crate) fn repairs(store: &dyn Store, folded: &[Folded]) -> Result<Vec<Repair>, SessionError> {
363 let torn = torn_tails(store, folded)?.into_iter();
364 let mut repairs: Vec<Repair> = torn
365 .map(|(log, kept, cut)| Repair::TornTail {
366 log,
367 kept,
368 cut,
369 set_aside: None,
370 })
371 .collect();
372 for log in folded {
373 repairs.extend(log.lost.iter().map(|why| Repair::LostImage {
374 log: log.name.clone(),
375 why: why.to_string(),
376 }));
377 }
378 Ok(repairs)
379}
380
381pub(crate) fn torn_tails(
386 store: &dyn Store,
387 folded: &[Folded],
388) -> Result<Vec<(String, u64, u64)>, SessionError> {
389 let mut torn = Vec::new();
390 for log in store.logs()? {
391 match folded.iter().find(|folded| folded.name == log) {
392 Some(folded) => torn.extend(folded.torn.map(|(kept, cut)| (log, kept, cut))),
393 None => torn.extend(torn_whole(store, &log).map(|cut| (log, 0, cut))),
394 }
395 }
396 Ok(torn)
397}
398
399fn torn_whole(store: &dyn Store, log: &str) -> Option<u64> {
402 let lines = log_lines(store, log).ok().flatten();
404 let lines = lines.filter(|lines| lines.len() == 0)?;
405 let cut = lines.tail().len() as u64;
406 (cut > 0).then_some(cut)
407}
408
409pub(crate) fn retired(logs: &[Folded]) -> HashSet<String> {
412 let logs = logs.iter().map(|log| {
413 let parent = log.header.parent.as_deref();
414 (log.name.as_str(), parent, log.retired)
415 });
416 skipped(logs).into_iter().map(str::to_owned).collect()
417}
418
419pub(crate) fn skipped<'a>(
422 logs: impl Iterator<Item = (&'a str, Option<&'a str>, bool)>,
423) -> HashSet<&'a str> {
424 let logs: Vec<_> = logs.collect();
425 let parents: HashMap<&str, Option<&str>> = logs
426 .iter()
427 .map(|&(name, parent, _)| (name, parent))
428 .collect();
429 let gone: HashSet<&str> = logs
430 .iter()
431 .filter(|&&(_, _, retired)| retired)
432 .map(|&(name, ..)| name)
433 .collect();
434 let mut skipped = HashSet::new();
435 for &(name, ..) in &logs {
436 let up = std::iter::successors(Some(name), |id| parents.get(id).copied().flatten());
438 if up.take(logs.len() + 1).any(|id| gone.contains(id)) {
439 skipped.insert(name);
440 }
441 }
442 skipped
443}
444
445pub(crate) fn session_of(folded: Vec<Folded>, repairs: Vec<Repair>) -> Session {
447 let mut session = Session {
448 repairs,
449 ..Session::default()
450 };
451 for log in folded {
452 if log.name == SESSION {
453 session.saved = log.state;
454 continue;
455 }
456 session.agents.push(SavedAgent {
457 id: log.name,
458 parent: log.header.parent,
459 started_ms: log.started_ms,
460 conversation: log.conversation,
461 saved: log.state,
462 retired: log.retired,
463 rewrites: log.header.rewrites,
464 });
465 }
466 session
467 .agents
468 .sort_by(|a, b| (a.started_ms, &a.id).cmp(&(b.started_ms, &b.id)));
469 session
470}
471
472impl Session {
473 pub fn read(store: &dyn Store) -> Result<Self, SessionError> {
482 let folded = fold_all(store, &read_logs(store)?)?;
483 let repairs = repairs(store, &folded)?;
484 Ok(session_of(folded, repairs))
485 }
486
487 pub fn read_lenient(store: &dyn Store) -> Result<(Self, Vec<Unreadable>), SessionError> {
497 let mut folded = Vec::new();
498 let mut unreadable = Vec::new();
499 for log in store.logs()? {
500 match fold(store, &log, Fold::WHOLE) {
501 Ok(read) => folded.extend(read),
502 Err(error) => unreadable.push(Unreadable { log, error }),
503 }
504 }
505 let repairs = repairs(store, &folded)?;
506 Ok((session_of(folded, repairs), unreadable))
507 }
508
509 pub fn agent(&self, id: &str) -> Option<&SavedAgent> {
511 self.agents.iter().find(|agent| agent.id == id)
512 }
513
514 pub fn agent_mut(&mut self, id: &str) -> Option<&mut SavedAgent> {
516 self.agents.iter_mut().find(|agent| agent.id == id)
517 }
518
519 pub fn save(&mut self, store: &dyn Store) -> Result<(), SessionError> {
536 let mut folded = Vec::new();
537 for (name, lines) in &read_logs(store)? {
538 folded.extend(tally(name, lines)?);
539 }
540 let had = |id: &str| folded.iter().find(|log| log.name == id);
541 for agent in &self.agents {
542 if agent.id == SESSION {
543 let why = format!("{SESSION:?} is the session's own log, not an agent's");
544 return Err(io::Error::new(io::ErrorKind::InvalidInput, why).into());
545 }
546 check_continues(had(&agent.id), &agent.id, &agent.conversation)?;
547 }
548 let mut book = Book::after(&folded);
549 let torn = torn_tails(store, &folded)?.into_iter();
550 book.set_aside(torn.map(|(log, _, cut)| (log, cut)));
551 book.start();
552 for agent in &mut self.agents {
553 let _ = book.opened(&agent.id, agent.parent.as_deref());
554 let put = &mut |name: String, bytes: Vec<u8>| put_blob(store, &name, &bytes);
555 book.save(put, &agent.id, &mut agent.conversation);
556 let known = had(&agent.id).map(|log| &log.state);
557 book.states(&agent.id, known, &agent.saved);
558 if agent.retired && !had(&agent.id).is_some_and(|log| log.retired) {
559 book.retire(&agent.id);
560 }
561 }
562 let known = had(SESSION).map(|log| &log.state);
563 book.states(SESSION, known, &self.saved);
564 book.flush(store)
565 }
566}
567
568fn check_continues(
572 log: Option<&Folded>,
573 name: &str,
574 conversation: &Conversation,
575) -> Result<(), SessionError> {
576 let Some(log) = log else {
577 return Ok(());
578 };
579 let changes = log.conversation.revision();
580 if changes > 0 && !conversation.continues(log.header.lineage, changes) {
581 return Err(SessionError::Changed {
582 log: name.to_owned(),
583 });
584 }
585 Ok(())
586}
587
588pub fn load(store: &dyn Store, log: &str) -> Result<Conversation, SessionError> {
602 check_leaf(log)?;
603 conversation_of(fold(store, log, Fold::WHOLE)?)
604}
605
606fn conversation_of(folded: Option<Folded>) -> Result<Conversation, SessionError> {
610 let Some(folded) = folded else {
611 return Ok(Conversation::new());
612 };
613 let mut lost = folded.lost.into_iter();
614 match lost.find(|failure| failure.kind() != io::ErrorKind::NotFound) {
615 Some(failure) => Err(failure.into()),
616 None => Ok(folded.conversation),
617 }
618}
619
620pub fn save(
650 store: &dyn Store,
651 log: &str,
652 conversation: &mut Conversation,
653) -> Result<(), SessionError> {
654 now(save_async(store, log, conversation))
655}
656
657pub async fn load_async<S: crate::AsyncStore + ?Sized>(
664 store: &S,
665 log: &str,
666) -> Result<Conversation, SessionError> {
667 check_leaf(log)?;
668 let bytes = match store.read_log_async(log).await {
669 Err(failure) if failure.kind() == io::ErrorKind::NotFound => {
670 return Ok(Conversation::new());
671 }
672 bytes => bytes?,
673 };
674 let mut named = std::collections::BTreeSet::new();
675 named_blobs(&bytes, &mut named);
676 let mut blobs = BTreeMap::new();
677 for name in named {
678 let read = store.read_file_async(&format!("{BLOBS}{name}")).await;
679 blobs.insert(name, read);
680 }
681 let blob = |name: &str| match blobs.get(name) {
682 Some(Ok(bytes)) => Ok(bytes.clone()),
683 Some(Err(failure)) => Err(io::Error::new(failure.kind(), failure.to_string())),
684 None => Err(io::ErrorKind::NotFound.into()),
685 };
686 conversation_of(fold_lines(
687 log,
688 &LogLines::parse(bytes),
689 Fold::WHOLE,
690 &blob,
691 )?)
692}
693
694pub async fn save_async<S: crate::AsyncStore + ?Sized>(
722 store: &S,
723 log: &str,
724 conversation: &mut Conversation,
725) -> Result<(), SessionError> {
726 let lines = match store.read_log_async(log).await {
727 Err(failure) if failure.kind() == io::ErrorKind::NotFound => LogLines::parse(Vec::new()),
728 bytes => LogLines::parse(bytes?),
729 };
730 let folded = tally(log, &lines)?;
731 check_continues(folded.as_ref(), log, conversation)?;
732 let mut book = Book::after(folded.as_slice());
733 let cut = lines.tail().len() as u64;
736 book.set_aside((cut > 0).then(|| (log.to_owned(), cut)));
737 let mut staged = conversation.clone();
738 book.start();
739 let _ = book.opened(log, None);
740 let mut blobs = Vec::new();
741 let put = &mut |name, bytes| {
742 blobs.push((name, bytes));
743 Ok(())
744 };
745 book.save(put, log, &mut staged);
746 for (name, bytes) in blobs {
747 store
748 .write_file_async(&format!("{BLOBS}{name}"), &bytes)
749 .await?;
750 }
751 book.flush_async(store).await?;
752 *conversation = staged;
753 Ok(())
754}
755
756pub fn history(store: &dyn Store, log: &str) -> Result<Vec<(Stamp, LogRecord)>, SessionError> {
764 let mut records = Vec::new();
765 read(store, log, |stamp, body| {
766 let record = match body {
767 Body::Change(mut recorded) => {
768 let _ = messages_mut(&mut recorded.change, &mut |message| {
770 let _ = load_images(message, store);
771 Ok(())
772 });
773 LogRecord::Change(recorded)
774 }
775 Body::Meta(Meta::State { key, value }) => LogRecord::Saved {
776 key,
777 value: value.filter(|value| !value.is_null()),
778 },
779 Body::Meta(Meta::Retired) => LogRecord::Retired,
780 Body::Meta(Meta::Header(_)) => return Ok(ControlFlow::Continue(())),
781 };
782 records.push((stamp, record));
783 Ok(ControlFlow::Continue(()))
784 })?;
785 Ok(records)
786}
787
788pub fn as_of(store: &dyn Store, log: &str, n: u64) -> Result<Conversation, SessionError> {
796 let how = Fold {
797 until: Some(n),
798 ..Fold::WHOLE
799 };
800 conversation_of(fold(store, log, how)?)
801}
802#[cfg(test)]
803mod tests;