rig_memory/
format.rs

1//! The saved format, [`FORMAT`]: one JSON-lines log per agent, named
2//! by the agent's id, plus the session's own ([`SESSION`]). Each line is a
3//! [`Stamp`] and one record: a `header` first, then `state` (a saved value
4//! by key; the latest wins, `null` removes it), `retired`, or one
5//! [`Change`](rig_core::conversation::Change) of the agent's conversation
6//! (`push`, `splice`, `halt`, `aside`, `reset`) with who made it (`by`).
7//! The header's `lineage` is the log's own
8//! [`Lineage`](rig_core::conversation::Lineage), chosen when it was
9//! started, so a conversation saved to it can be told from one that only
10//! has as many changes. Its `rewrites` say who rewrote the log in place
11//! ([`rewrite`](crate::rewrite)) and when; each rewrite gives the log a
12//! fresh lineage.
13//!
14//! A record keeps when each entry was made: a new entry was made when its
15//! record says (`t`), an entry put back as it was says so itself (`made`),
16//! and a reset's state carries each entry's.
17//!
18//! Every string a line holds is redacted as the line is written
19//! ([`rig_core::redact::json_text`]), so no API key or token lands in a log.
20//!
21//! ```json
22//! {"seq":0,"n":0,"t":1791779409181,"type":"header","format":"rig-session/4","log":"cd9a…","lineage":8125…}
23//! {"seq":1,"n":1,"t":1791779409181,"type":"state","key":"rig_ecs::agent::ToolAccess","value":"All"}
24//! {"seq":2,"n":2,"t":1791779409182,"type":"push","entries":[{"message":{…},"id":0}]}
25//! ```
26
27use std::borrow::Cow;
28use std::fmt;
29use std::io;
30
31use rig_core::conversation::{
32    Aside, AsideId, Author, Change, ChangeRecord, Conversation, Entry, Halt, Lineage, MessageId,
33};
34use rig_core::redact;
35use serde::{Deserialize, Serialize};
36use serde_json::Value;
37
38/// The saved format a log's header names: the one format this rig reads.
39pub const FORMAT: &str = "rig-session/4";
40
41/// The log of the session's own saved values, beside the agents' logs.
42pub const SESSION: &str = "session";
43
44/// Where a record is in a session: its place among the session's records
45/// in the order they were made (`n`, what [`as_of`](crate::as_of) cuts
46/// by), its place in its log (`seq`, dense from
47/// the header's 0), and when it was written (`t`, milliseconds since the
48/// Unix epoch).
49#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
50pub struct Stamp {
51    /// Its place among the session's records: it increases along every
52    /// log, and across the logs one [`Writer`](crate::Writer) or
53    /// [`Session::save`](crate::Session::save) wrote.
54    pub n: u64,
55    /// Its place in its log.
56    pub seq: u64,
57    /// When it was written.
58    pub t: u64,
59}
60
61/// Why a session cannot be read or saved. Reading is strict: a session
62/// with any record it cannot read is refused whole, and nothing is written.
63#[derive(Debug, thiserror::Error)]
64pub enum SessionError {
65    /// The store failed.
66    #[error(transparent)]
67    Io(#[from] io::Error),
68    /// The directory holds the single-file layout older than `rig-session/2` (a
69    /// top-level `*.jsonl`); it is left as it is.
70    #[error("this session was made by another version of rig; start a new session")]
71    OldLayout,
72    /// A log's header names a format other than [`FORMAT`], or none; it is
73    /// left as it is.
74    #[error(
75        "this session was made by another version of rig; start a new session (the log `{log}` \
76         is in format {}, not {FORMAT})",
77        .found.as_deref().unwrap_or("(none)")
78    )]
79    Format {
80        /// The log.
81        log: String,
82        /// The format its header names.
83        found: Option<String>,
84    },
85    /// A record cannot be read, or its change does not apply.
86    #[error("the log `{log}` cannot be read at line {line}: {why}")]
87    Record {
88        /// The log.
89        log: String,
90        /// The line, from 1.
91        line: usize,
92        /// Why.
93        why: String,
94    },
95    /// The log moved on since the conversation saved to it was read: saving
96    /// would lose the changes made there.
97    #[error("the log `{log}` changed since it was read")]
98    Changed {
99        /// The log.
100        log: String,
101    },
102}
103
104impl From<SessionError> for io::Error {
105    /// The store's own error back; any other as
106    /// [`io::ErrorKind::InvalidData`], so `?` works in code that returns
107    /// [`io::Result`].
108    fn from(error: SessionError) -> Self {
109        match error {
110            SessionError::Io(error) => error,
111            error => io::Error::new(io::ErrorKind::InvalidData, error),
112        }
113    }
114}
115
116/// What reading a session repaired, or found it must work around.
117#[derive(Clone, Debug, PartialEq, Eq)]
118pub enum Repair {
119    /// The log ended in a line a crash cut short. The log is read up to
120    /// `kept` bytes; the `cut` bytes after were set aside in the named file
121    /// `set_aside` when the log was opened for writing (`None` when it was
122    /// only read).
123    TornTail {
124        /// The log.
125        log: String,
126        /// The bytes read.
127        kept: u64,
128        /// The bytes cut.
129        cut: u64,
130        /// The file they were set aside in.
131        set_aside: Option<String>,
132    },
133    /// An image's blob cannot be read, so the image has no data.
134    LostImage {
135        /// The log.
136        log: String,
137        /// Why.
138        why: String,
139    },
140}
141
142impl fmt::Display for Repair {
143    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
144        match self {
145            Self::TornTail {
146                log,
147                cut,
148                set_aside,
149                ..
150            } => {
151                write!(f, "the log `{log}` ended in a line cut short ({cut} bytes)")?;
152                match set_aside {
153                    Some(file) => write!(f, ", set aside in {file}"),
154                    None => Ok(()),
155                }
156            }
157            Self::LostImage { log, why } => write!(f, "an image of the log `{log}` is lost: {why}"),
158        }
159    }
160}
161
162/// The first record of a log.
163#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
164pub(crate) struct Header {
165    #[serde(default, skip_serializing_if = "Option::is_none")]
166    pub format: Option<String>,
167    #[serde(default)]
168    pub log: String,
169    /// The log of the agent that spawned this one.
170    #[serde(default, skip_serializing_if = "Option::is_none")]
171    pub parent: Option<String>,
172    /// The log's own lineage, which every conversation loaded from it or
173    /// saved to it carries.
174    #[serde(default)]
175    pub lineage: Lineage,
176    /// Who rewrote the log in place, and when, oldest first.
177    #[serde(default, skip_serializing_if = "Vec::is_empty")]
178    pub rewrites: Vec<LogRewrite>,
179}
180
181/// A rewrite of a log in place ([`rewrite`](crate::rewrite)): when, and by
182/// whom, such as the plugin that redacted it.
183#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
184pub struct LogRewrite {
185    /// When, in milliseconds since the Unix epoch.
186    pub t: u64,
187    /// Who.
188    pub by: Author,
189}
190
191impl Header {
192    /// The header of a log just started, with a fresh lineage.
193    pub fn new(log: &str, parent: Option<&str>) -> Self {
194        Self {
195            format: Some(FORMAT.to_owned()),
196            log: log.to_owned(),
197            parent: parent.map(str::to_owned),
198            lineage: Lineage::fresh(),
199            rewrites: Vec::new(),
200        }
201    }
202}
203
204/// A record that is no change of the conversation.
205#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
206#[serde(tag = "type", rename_all = "snake_case")]
207pub(crate) enum Meta {
208    Header(Header),
209    State {
210        key: String,
211        #[serde(default)]
212        value: Option<Value>,
213    },
214    Retired,
215}
216
217/// A record as read.
218pub(crate) enum Body {
219    Meta(Meta),
220    /// Its `at` is the stamp's `t`.
221    Change(ChangeRecord),
222}
223
224/// The fields every line starts with.
225#[derive(Deserialize)]
226struct Head<'a> {
227    seq: u64,
228    n: u64,
229    t: u64,
230    #[serde(rename = "type", borrow)]
231    kind: std::borrow::Cow<'a, str>,
232}
233
234/// The record types that are changes of the conversation.
235const CHANGES: [&str; 5] = ["push", "splice", "halt", "aside", "reset"];
236
237/// A line's stamp, and whether it is a change of the conversation, read
238/// without the rest of it.
239pub(crate) fn stamp_of(line: &[u8]) -> Result<(Stamp, bool), String> {
240    let head: Head<'_> = serde_json::from_slice(line).map_err(|error| error.to_string())?;
241    Ok((head.stamp(), CHANGES.contains(&&*head.kind)))
242}
243
244impl Head<'_> {
245    fn stamp(&self) -> Stamp {
246        Stamp {
247            n: self.n,
248            seq: self.seq,
249            t: self.t,
250        }
251    }
252}
253
254/// A line read in one pass: its stamp and type, and, for a change, who
255/// made it and every field a change has, of which the type says which it
256/// uses. A change read whole without its stamp would be buffered twice:
257/// once for the fields around it, once for its type.
258#[derive(Deserialize)]
259struct Read<'a> {
260    seq: u64,
261    n: u64,
262    t: u64,
263    #[serde(rename = "type", borrow)]
264    kind: std::borrow::Cow<'a, str>,
265    #[serde(default)]
266    by: Option<Author>,
267    entries: Option<Vec<Entry>>,
268    from: Option<usize>,
269    to: Option<usize>,
270    removed: Option<Vec<MessageId>>,
271    with: Option<Vec<Entry>>,
272    #[serde(default)]
273    halt: Option<Halt>,
274    id: Option<AsideId>,
275    #[serde(default)]
276    aside: Option<Aside>,
277    state: Option<Box<Conversation>>,
278}
279
280impl Read<'_> {
281    /// The change the line holds.
282    fn change(self) -> Result<Change, String> {
283        fn given<T>(field: Option<T>, name: &str) -> Result<T, String> {
284            field.ok_or_else(|| format!("missing field `{name}`"))
285        }
286        Ok(match &*self.kind {
287            "push" => Change::Push {
288                entries: given(self.entries, "entries")?,
289            },
290            "splice" => Change::Splice {
291                from: given(self.from, "from")?,
292                to: given(self.to, "to")?,
293                removed: given(self.removed, "removed")?,
294                with: given(self.with, "with")?,
295            },
296            "halt" => Change::Halt { halt: self.halt },
297            "aside" => Change::Aside {
298                id: given(self.id, "id")?,
299                aside: self.aside,
300            },
301            "reset" => Change::Reset {
302                state: given(self.state, "state")?,
303            },
304            other => return Err(format!("unknown record type {other:?}")),
305        })
306    }
307}
308
309/// Reads one line.
310pub(crate) fn parse(line: &[u8]) -> Result<(Stamp, Body), String> {
311    let mut read: Read<'_> = serde_json::from_slice(line).map_err(|error| error.to_string())?;
312    let stamp = Stamp {
313        n: read.n,
314        seq: read.seq,
315        t: read.t,
316    };
317    let body = match &*read.kind {
318        "header" | "state" | "retired" => {
319            Body::Meta(serde_json::from_slice(line).map_err(|error| error.to_string())?)
320        }
321        _ => {
322            let by = read.by.take();
323            Body::Change(ChangeRecord {
324                change: read.change()?,
325                by,
326                at: Some(stamp.t),
327            })
328        }
329    };
330    Ok((stamp, body))
331}
332
333/// A line as written.
334#[derive(Serialize)]
335struct Line<'a> {
336    seq: u64,
337    n: u64,
338    t: u64,
339    #[serde(skip_serializing_if = "Option::is_none")]
340    by: Option<&'a Author>,
341    #[serde(flatten)]
342    body: Out<'a>,
343}
344
345#[derive(Serialize)]
346#[serde(untagged)]
347enum Out<'a> {
348    Meta(&'a Meta),
349    Change(&'a Change),
350}
351
352/// Writes a meta record line, newline included, to `out`.
353pub(crate) fn write_meta(out: &mut Vec<u8>, stamp: Stamp, meta: &Meta) -> serde_json::Result<()> {
354    write(out, stamp, None, Out::Meta(meta))
355}
356
357/// Writes a change line, newline included, to `out`.
358pub(crate) fn write_change(
359    out: &mut Vec<u8>,
360    stamp: Stamp,
361    by: Option<&Author>,
362    change: &Change,
363) -> serde_json::Result<()> {
364    write(out, stamp, by, Out::Change(change))
365}
366
367/// Writes a line, every secret in its strings redacted
368/// ([`redact::json_text`]), so no key a conversation or a saved value holds
369/// lands in a log.
370fn write(
371    out: &mut Vec<u8>,
372    stamp: Stamp,
373    by: Option<&Author>,
374    body: Out<'_>,
375) -> serde_json::Result<()> {
376    let line = Line {
377        seq: stamp.seq,
378        n: stamp.n,
379        t: stamp.t,
380        by,
381        body,
382    };
383    let start = out.len();
384    serde_json::to_writer(&mut *out, &line)?;
385    // No key reaches the disk: every string of the line is redacted. JSON
386    // is UTF-8, and a line without a secret is not copied.
387    let written = out
388        .get(start..)
389        .and_then(|line| std::str::from_utf8(line).ok());
390    if let Some(Cow::Owned(redacted)) = written.map(redact::json_text) {
391        out.truncate(start);
392        out.extend_from_slice(redacted.as_bytes());
393    }
394    out.push(b'\n');
395    Ok(())
396}
397
398/// The complete lines of a log as stored. A torn tail a crash left (bytes
399/// after the last newline, or a last line that is not JSON) is not one.
400#[derive(Clone, Debug, Default)]
401pub(crate) struct LogLines {
402    bytes: Vec<u8>,
403    /// Where each complete line, newline included, ends.
404    ends: Vec<usize>,
405}
406
407impl LogLines {
408    pub fn parse(bytes: Vec<u8>) -> Self {
409        let mut ends: Vec<usize> = memchr::memchr_iter(b'\n', &bytes)
410            .map(|at| at + 1)
411            .collect();
412        let last_start = ends
413            .len()
414            .checked_sub(2)
415            .and_then(|at| ends.get(at))
416            .copied();
417        let last = ends
418            .last()
419            .and_then(|end| bytes.get(last_start.unwrap_or(0)..end.saturating_sub(1)));
420        if last.is_some_and(|line| {
421            !line.is_empty() && serde_json::from_slice::<serde::de::IgnoredAny>(line).is_err()
422        }) {
423            ends.pop();
424        }
425        Self { bytes, ends }
426    }
427
428    /// Each complete line that is not blank, without its newline, with its
429    /// number from 1 and the byte offset just past it.
430    pub fn iter(&self) -> impl Iterator<Item = (usize, u64, &[u8])> {
431        let starts = std::iter::once(0).chain(self.ends.iter().copied());
432        starts
433            .zip(self.ends.iter().copied())
434            .enumerate()
435            .filter_map(|(at, (start, end))| {
436                let line = self.bytes.get(start..end.saturating_sub(1))?;
437                (!line.is_empty()).then_some((at + 1, end as u64, line))
438            })
439    }
440
441    /// The bytes the complete lines cover.
442    pub fn len(&self) -> u64 {
443        self.ends.last().map_or(0, |end| *end as u64)
444    }
445
446    /// The bytes after the complete lines: a torn tail.
447    pub fn tail(&self) -> &[u8] {
448        let len = usize::try_from(self.len()).unwrap_or(usize::MAX);
449        self.bytes.get(len..).unwrap_or_default()
450    }
451
452    /// The log's first `len` bytes, no further than its complete lines.
453    pub fn prefix(&self, len: u64) -> &[u8] {
454        let len = usize::try_from(len.min(self.len())).unwrap_or(0);
455        self.bytes.get(..len).unwrap_or_default()
456    }
457}
458
459#[cfg(test)]
460mod tests;
rig_memory/format.rs

rig_memory/
format.rs
