repositories / smith
smith
There are many coding harnesses - but this one is fast
owned by admin
smith-core/src/session.rs
Raw//! Branching session state: an index arena of fixed-size entry records over
//! one contiguous payload arena, plus restoration from durable frames
//! (`SMH-SPEC-SPEC0001`, Sessions).
//!
//! Entry content lives as canonical CBOR bytes, at rest and in memory. A
//! record carries what traversal needs ([`EntryKind`]); anything else is a
//! decode into the consuming scope.
use crate::frame::Frame;
use serde::{Deserialize, Serialize};
use smith::error::{Result, SmithError};
use smith::id::{BranchId, EntryId, MessageId, SessionId, ToolCallId};
use smith::message::{ContentBlock, Message, Role};
use std::collections::HashMap;
/// Meta entry kind that marks a compaction boundary.
pub const COMPACTION_KIND: &str = "compaction";
/// Metadata kind for one redacted provider HTTP effect record.
pub const PROVIDER_HTTP_KIND: &str = "provider-http";
/// Deterministic JSON: recursively sort object keys so equivalent values
/// encode to identical CBOR regardless of insertion order.
fn canonical_json(value: &serde_json::Value) -> serde_json::Value {
match value {
serde_json::Value::Object(map) => {
let mut keys: Vec<&String> = map.keys().collect();
keys.sort();
let mut out = serde_json::Map::new();
for key in keys {
out.insert(key.clone(), canonical_json(&map[key]));
}
serde_json::Value::Object(out)
}
serde_json::Value::Array(items) => {
serde_json::Value::Array(items.iter().map(canonical_json).collect())
}
other => other.clone(),
}
}
/// Whether `value` already has the key order [`canonical_json`] produces.
fn is_canonical(value: &serde_json::Value) -> bool {
match value {
serde_json::Value::Object(map) => {
map.keys().zip(map.keys().skip(1)).all(|(a, b)| a < b) && map.values().all(is_canonical)
}
serde_json::Value::Array(items) => items.iter().all(is_canonical),
_ => true,
}
}
/// Durable session identity, persisted as the first frame of a session file
/// so reload reproduces the same IDs instead of minting new ones.
#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct SessionHeader {
/// Session identifier.
pub session_id: SessionId,
/// Identifier of the root branch.
pub branch_id: BranchId,
}
/// Byte range of one entry's content in its session's payload arena.
///
/// `u32` offsets cap one session arena at 4 GiB.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Span {
/// First content byte.
pub offset: u32,
/// Content length in bytes.
pub len: u32,
}
impl Span {
/// The span of `len` bytes at `offset`, or an error past the arena cap.
fn at(offset: usize, len: usize) -> Result<Self> {
let full = || SmithError::Session {
code: "SESSION_ARENA_FULL".to_string(),
message: format!("entry content at {offset}+{len} exceeds the 4 GiB session arena"),
};
let end = u32::try_from(offset.saturating_add(len)).map_err(|_| full())?;
let offset = u32::try_from(offset).map_err(|_| full())?;
Ok(Self {
offset,
len: end - offset,
})
}
const fn end(self) -> usize {
self.offset as usize + self.len as usize
}
const fn range(self) -> std::ops::Range<usize> {
self.offset as usize..self.end()
}
}
/// What an entry recorded, reduced to the identifiers traversal pairs on.
#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub enum EntryKind {
/// Chat message.
Message {
/// Message identifier.
id: MessageId,
/// Who produced the message.
role: Role,
},
/// Tool call.
ToolCall {
/// Identifier pairing the call with its result.
call_id: ToolCallId,
},
/// Tool result.
ToolResult {
/// Identifier of the call this result answers.
call_id: ToolCallId,
/// Whether the effect succeeded.
ok: bool,
},
/// Session metadata.
Meta {
/// Whether the entry marks a compaction boundary.
compaction: bool,
},
}
/// A fixed-size session entry record; its content lives in the session
/// payload arena.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Entry {
/// Entry ID.
pub id: EntryId,
/// Parent entry ID (None for root).
pub parent: Option<EntryId>,
/// Timestamp (Unix ms).
pub timestamp_ms: i64,
/// What the entry recorded.
pub kind: EntryKind,
/// Canonical CBOR of the [`EntryContent`], in the session arena.
pub content: Span,
}
impl Entry {
/// Whether this entry marks a compaction boundary.
#[must_use]
pub const fn is_compaction(&self) -> bool {
matches!(self.kind, EntryKind::Meta { compaction: true })
}
}
/// Wire shape of one entry frame: the record plus its content as an opaque
/// canonical CBOR byte string.
#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct EntryFrame {
/// Entry ID.
pub id: EntryId,
/// Parent entry ID (None for root).
pub parent: Option<EntryId>,
/// Timestamp (Unix ms).
pub timestamp_ms: i64,
/// What the entry recorded.
pub kind: EntryKind,
/// Canonical CBOR of the [`EntryContent`].
#[serde(serialize_with = "cbor_bytes", deserialize_with = "byte_buf")]
pub content: Vec<u8>,
}
impl EntryFrame {
/// A frame for `content` with a fresh identity, outside any session.
///
/// # Errors
///
/// Returns [`SmithError::Session`] when the content cannot be encoded.
pub fn new(parent: Option<EntryId>, content: &EntryContent) -> Result<Self> {
let mut bytes = Vec::new();
content.encode_into(&mut bytes)?;
Ok(Self {
id: EntryId::new(),
parent,
timestamp_ms: now_ms(),
kind: content.kind(),
content: bytes,
})
}
}
/// Borrowed [`EntryFrame`] serialization of a session record and its arena
/// bytes, so arena content is framed without being copied out first.
#[derive(Serialize)]
pub(crate) struct EntryWire<'a> {
id: EntryId,
parent: Option<EntryId>,
timestamp_ms: i64,
kind: EntryKind,
#[serde(serialize_with = "cbor_bytes")]
content: &'a [u8],
}
impl<'a> EntryWire<'a> {
pub(crate) const fn of(entry: &Entry, content: &'a [u8]) -> Self {
Self {
id: entry.id,
parent: entry.parent,
timestamp_ms: entry.timestamp_ms,
kind: entry.kind,
content,
}
}
}
/// Serialize bytes as one CBOR byte string rather than an integer array.
fn cbor_bytes<S: serde::Serializer>(
bytes: &impl AsRef<[u8]>,
serializer: S,
) -> std::result::Result<S::Ok, S::Error> {
serializer.serialize_bytes(bytes.as_ref())
}
fn byte_buf<'de, D: serde::Deserializer<'de>>(
deserializer: D,
) -> std::result::Result<Vec<u8>, D::Error> {
struct Visitor;
impl serde::de::Visitor<'_> for Visitor {
type Value = Vec<u8>;
fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("a byte string")
}
fn visit_bytes<E: serde::de::Error>(self, bytes: &[u8]) -> std::result::Result<Vec<u8>, E> {
Ok(bytes.to_vec())
}
fn visit_byte_buf<E: serde::de::Error>(
self,
bytes: Vec<u8>,
) -> std::result::Result<Vec<u8>, E> {
Ok(bytes)
}
}
deserializer.deserialize_byte_buf(Visitor)
}
/// What one entry recorded; the decoded view of an entry's content.
#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
#[serde(tag = "kind", content = "data")]
pub enum EntryContent {
/// Chat message.
Message(Message),
/// Tool call (recorded for replay).
ToolCall {
/// Invoked tool.
name: String,
/// Canonical tool arguments.
input: serde_json::Value,
/// Identifier pairing this call with its result.
call_id: ToolCallId,
},
/// Tool result.
ToolResult {
/// Whether the effect succeeded.
ok: bool,
/// Recorded output or error text.
output: String,
/// Identifier of the call this result answers.
call_id: ToolCallId,
},
/// Session metadata event (compaction, branch selection).
Meta {
/// Metadata kind, for example [`COMPACTION_KIND`].
kind: String,
/// Kind-specific payload.
detail: serde_json::Value,
},
}
impl EntryContent {
/// The record kind of this content.
#[must_use]
pub fn kind(&self) -> EntryKind {
EntryContentRef::from(self).kind()
}
/// Decode content from its canonical CBOR bytes.
///
/// # Errors
///
/// Returns [`SmithError::Session`] when `bytes` are not entry content.
pub fn decode(bytes: &[u8]) -> Result<Self> {
ciborium::de::from_reader(bytes).map_err(|e| SmithError::Session {
code: "ENTRY_DECODE".to_string(),
message: e.to_string(),
})
}
/// Append the canonical CBOR of this content to `out`.
fn encode_into(&self, out: &mut Vec<u8>) -> Result<()> {
EntryContentRef::from(self).encode_into(out)
}
}
/// Borrowed [`EntryContent`]: what recording needs without owning a copy.
///
/// The serialize shape is the same as [`EntryContent`], so both produce the
/// same canonical bytes; this is the one place content is canonicalized.
#[derive(Clone, Copy, Debug, Serialize)]
#[serde(tag = "kind", content = "data")]
pub enum EntryContentRef<'a> {
/// Chat message.
Message(&'a Message),
/// Tool call.
ToolCall {
/// Invoked tool.
name: &'a str,
/// Tool arguments.
input: &'a serde_json::Value,
/// Identifier pairing this call with its result.
call_id: ToolCallId,
},
/// Tool result.
ToolResult {
/// Whether the effect succeeded.
ok: bool,
/// Output or error text.
output: &'a str,
/// Identifier of the call this result answers.
call_id: ToolCallId,
},
/// Session metadata event.
Meta {
/// Metadata kind.
kind: &'a str,
/// Kind-specific payload.
detail: &'a serde_json::Value,
},
}
impl<'a> From<&'a EntryContent> for EntryContentRef<'a> {
fn from(content: &'a EntryContent) -> Self {
match content {
EntryContent::Message(message) => Self::Message(message),
EntryContent::ToolCall {
name,
input,
call_id,
} => Self::ToolCall {
name,
input,
call_id: *call_id,
},
EntryContent::ToolResult {
ok,
output,
call_id,
} => Self::ToolResult {
ok: *ok,
output,
call_id: *call_id,
},
EntryContent::Meta { kind, detail } => Self::Meta { kind, detail },
}
}
}
impl EntryContentRef<'_> {
/// The record kind of this content.
#[must_use]
pub fn kind(&self) -> EntryKind {
match *self {
Self::Message(message) => EntryKind::Message {
id: message.id,
role: message.role,
},
Self::ToolCall { call_id, .. } => EntryKind::ToolCall { call_id },
Self::ToolResult { call_id, ok, .. } => EntryKind::ToolResult { call_id, ok },
Self::Meta { kind, .. } => EntryKind::Meta {
compaction: kind == COMPACTION_KIND,
},
}
}
/// Append the canonical CBOR of this content to `out`; JSON values are
/// re-ordered into canonical key order only when they are not already.
fn encode_into(&self, out: &mut Vec<u8>) -> Result<()> {
let canonical_value: serde_json::Value;
let canonical_message: Message;
// The view borrows from `self` or from a local canonical copy.
let view: EntryContentRef<'_> = match *self {
Self::ToolCall {
name,
input,
call_id,
} if !is_canonical(input) => {
canonical_value = canonical_json(input);
EntryContentRef::ToolCall {
name,
input: &canonical_value,
call_id,
}
}
Self::Meta { kind, detail } if !is_canonical(detail) => {
canonical_value = canonical_json(detail);
EntryContentRef::Meta {
kind,
detail: &canonical_value,
}
}
Self::Message(message) if !message_is_canonical(message) => {
canonical_message = canonical_message_of(message);
EntryContentRef::Message(&canonical_message)
}
other => other,
};
ciborium::ser::into_writer(&view, out).map_err(|e| SmithError::Session {
code: "ENTRY_ENCODE".to_string(),
message: e.to_string(),
})
}
}
fn message_is_canonical(message: &Message) -> bool {
message.blocks.iter().all(|block| match block {
ContentBlock::ToolUse { input, .. } => is_canonical(input),
_ => true,
})
}
fn canonical_message_of(message: &Message) -> Message {
let mut message = message.clone();
for block in &mut message.blocks {
if let ContentBlock::ToolUse { input, .. } = block {
*input = canonical_json(input);
}
}
message
}
/// Wall-clock milliseconds since the Unix epoch, saturating on overflow.
fn now_ms() -> i64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |elapsed| {
i64::try_from(elapsed.as_millis()).unwrap_or(i64::MAX)
})
}
/// One root-to-leaf chain of a session.
///
/// Only its [`Session`] mutates a branch, so every record's span points into
/// that session's arena.
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct Branch {
id: BranchId,
entries: Vec<Entry>,
selected_idx: usize,
}
impl Branch {
const fn with_id(id: BranchId) -> Self {
Self {
id,
entries: Vec::new(),
selected_idx: 0,
}
}
/// Branch identity.
#[must_use]
pub const fn id(&self) -> BranchId {
self.id
}
/// Entry records ordered root to leaf.
#[must_use]
pub fn entries(&self) -> &[Entry] {
&self.entries
}
/// Append `entry` and select it.
fn append(&mut self, entry: Entry) -> &Entry {
self.selected_idx = self.entries.len();
self.entries.push(entry);
&self.entries[self.selected_idx]
}
/// The currently selected entry, if the branch has any.
#[must_use]
pub fn selected_entry(&self) -> Option<&Entry> {
self.entries.get(self.selected_idx)
}
}
/// A session is a persistent, branching message/tool history.
///
/// Fields stay private so the invariants hold: at least one branch, a valid
/// active index, every span inside the arena, and a compaction boundary
/// that names a recorded compaction entry.
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct Session {
id: SessionId,
/// All branches (default branch is index 0).
branches: Vec<Branch>,
/// Canonical CBOR content of every entry, addressed by [`Span`].
payloads: Vec<u8>,
active_branch_idx: usize,
compaction_boundary: Option<EntryId>,
}
impl Default for Session {
fn default() -> Self {
Self::new()
}
}
impl Session {
/// A new session with one empty branch.
#[must_use]
pub fn new() -> Self {
Self {
id: SessionId::new(),
branches: vec![Branch::with_id(BranchId::new())],
payloads: Vec::new(),
active_branch_idx: 0,
compaction_boundary: None,
}
}
/// Session identifier.
#[must_use]
pub const fn id(&self) -> SessionId {
self.id
}
/// All branches; the active one is [`Session::active_branch`].
#[must_use]
pub fn branches(&self) -> &[Branch] {
&self.branches
}
/// The latest compaction boundary entry, if any.
#[must_use]
pub const fn compaction_boundary(&self) -> Option<EntryId> {
self.compaction_boundary
}
/// The selected branch.
#[must_use]
pub fn active_branch(&self) -> &Branch {
&self.branches[self.active_branch_idx]
}
/// Identifier of the selected entry, if any.
#[must_use]
pub fn selected_id(&self) -> Option<EntryId> {
self.active_branch().selected_entry().map(|e| e.id)
}
/// Encode `content` once into the arena and append its record to the
/// selected branch, selecting it. A compaction entry becomes the
/// compaction boundary.
///
/// # Errors
///
/// Returns [`SmithError::Session`] when the content cannot be encoded or
/// the arena is full; the session is unchanged then.
pub fn append<'a>(&mut self, content: impl Into<EntryContentRef<'a>>) -> Result<&Entry> {
self.append_with(content, |_, _| Ok(()))
}
/// [`Session::append`], running `persist` on the new record and its
/// content bytes before the entry becomes part of the session.
///
/// # Errors
///
/// As [`Session::append`], or the error of `persist`; the session is
/// unchanged then.
pub fn append_with<'a>(
&mut self,
content: impl Into<EntryContentRef<'a>>,
persist: impl FnOnce(&Entry, &[u8]) -> Result<()>,
) -> Result<&Entry> {
let content = content.into();
let parent = self.selected_id();
let offset = self.payloads.len();
let staged = content
.encode_into(&mut self.payloads)
.and_then(|()| Span::at(offset, self.payloads.len() - offset))
.map(|span| Entry {
id: EntryId::new(),
parent,
timestamp_ms: now_ms(),
kind: content.kind(),
content: span,
})
.and_then(|entry| persist(&entry, &self.payloads[offset..]).map(|()| entry));
let entry = match staged {
Ok(entry) => entry,
Err(error) => {
self.payloads.truncate(offset);
return Err(error);
}
};
if entry.is_compaction() {
self.compaction_boundary = Some(entry.id);
}
Ok(self.branches[self.active_branch_idx].append(entry))
}
/// The canonical CBOR content bytes of `entry`; empty when `entry` does
/// not belong to this session.
#[must_use]
pub fn bytes(&self, entry: &Entry) -> &[u8] {
self.payloads.get(entry.content.range()).unwrap_or_default()
}
/// Decode the content of `entry` into the caller's scope, verifying it
/// against the record's [`EntryKind`], which traversal trusts unchecked.
///
/// # Errors
///
/// Returns [`SmithError::Session`] with code `ENTRY_DECODE` when the
/// bytes do not decode, or `ENTRY_KIND_MISMATCH` when they decode to
/// content of another kind than the record claims.
pub fn content(&self, entry: &Entry) -> Result<EntryContent> {
let content = EntryContent::decode(self.bytes(entry))?;
let found = content.kind();
if found != entry.kind {
return Err(SmithError::Session {
code: "ENTRY_KIND_MISMATCH".to_string(),
message: format!(
"entry {} records {:?} but holds {found:?}",
entry.id, entry.kind
),
});
}
Ok(content)
}
/// The root-to-`at` entry chain, or `None` when `at` is unknown or its
/// ancestry is broken.
///
/// A branch stores its root-to-leaf chain in order, so the chain is the
/// branch prefix ending at `at`; each parent link is verified against the
/// preceding entry and no lookup table is built.
fn ancestry(&self, at: EntryId) -> Option<&[Entry]> {
let (branch, idx) = self.branches.iter().find_map(|branch| {
let idx = branch.entries.iter().position(|entry| entry.id == at)?;
Some((branch, idx))
})?;
let chain = &branch.entries[..=idx];
let rooted = chain.first().is_some_and(|root| root.parent.is_none());
let linked = chain
.windows(2)
.all(|pair| pair[1].parent == Some(pair[0].id));
(rooted && linked).then_some(chain)
}
/// Fork the history ending at `at` into a new session.
///
/// Entry IDs, order, and timestamps are preserved, `at` becomes the
/// selected leaf, and the compaction boundary carries over when it lies
/// on the retained path. The new session gets fresh session and branch
/// identity, because it is a different durable history. The retained
/// records are copied with the one arena prefix their spans cover, so
/// spans stay valid unchanged.
///
/// # Errors
///
/// Returns [`SmithError::Session`] when `at` is not part of this session.
pub fn fork_at(&self, at: EntryId) -> Result<Self> {
let entries = self.ancestry(at).ok_or_else(|| SmithError::Session {
code: "ENTRY_NOT_FOUND".to_string(),
message: format!("entry {at} is not part of session {}", self.id),
})?;
let boundary = self
.compaction_boundary
.filter(|id| entries.iter().any(|entry| entry.id == *id));
let end = entries
.iter()
.map(|entry| entry.content.end())
.max()
.unwrap_or(0);
Ok(Self {
id: SessionId::new(),
branches: vec![Branch {
id: BranchId::new(),
entries: entries.to_vec(),
selected_idx: entries.len().saturating_sub(1),
}],
payloads: self.payloads.get(..end).unwrap_or_default().to_vec(),
active_branch_idx: 0,
compaction_boundary: boundary,
})
}
/// The durable identity of this session's active branch.
#[must_use]
pub fn header(&self) -> SessionHeader {
SessionHeader {
session_id: self.id,
branch_id: self.active_branch().id,
}
}
/// Restore in-memory state from decoded frames.
///
/// A header frame restores the session and branch identity, so state
/// rebuilt from the same file is identical on every load. Stored parent
/// links define the branches: each entry with no children is a leaf, and
/// its root-to-leaf chain is one branch. The last stored entry is the
/// selected leaf of the active branch. Unknown frames are skipped here,
/// and compaction entries restore the compaction boundary.
///
/// Content bytes are copied into the arena, never decoded; every
/// collection is sized up front, so a linear history restores with a
/// constant number of allocations.
///
/// # Errors
///
/// Returns [`SmithError::Session`] when the content exceeds the arena cap.
pub fn from_frames(frames: &[Frame]) -> Result<Self> {
let mut header = None;
let (mut count, mut total) = (0, 0);
for frame in frames {
match frame {
Frame::Header { header: found } => header = Some(*found),
Frame::Known { entry } => {
count += 1;
total += entry.content.len();
}
Frame::Unknown { .. } => {}
}
}
Span::at(0, total)?;
let mut payloads = Vec::with_capacity(total);
let mut stored = Vec::with_capacity(count);
let mut boundary = None;
for frame in frames {
let Frame::Known { entry } = frame else {
continue;
};
let record = Entry {
id: entry.id,
parent: entry.parent,
timestamp_ms: entry.timestamp_ms,
kind: entry.kind,
content: Span::at(payloads.len(), entry.content.len())?,
};
payloads.extend_from_slice(&entry.content);
if record.is_compaction() {
boundary = Some(record.id);
}
stored.push(record);
}
let header = header.unwrap_or_else(|| SessionHeader {
session_id: SessionId::new(),
branch_id: BranchId::new(),
});
let (branches, active_branch_idx) = restore_branches(&stored, header.branch_id);
Ok(Self {
id: header.session_id,
branches,
payloads,
active_branch_idx,
compaction_boundary: boundary,
})
}
}
/// Rebuild branches from stored records and their parent links, returning
/// them with the index of the active branch.
fn restore_branches(stored: &[Entry], active_id: BranchId) -> (Vec<Branch>, usize) {
let Some(last) = stored.last() else {
return (vec![Branch::with_id(active_id)], 0);
};
// Entry id to its index and whether any stored entry names it as parent.
let mut nodes: HashMap<EntryId, (usize, bool)> = HashMap::with_capacity(stored.len());
for (idx, entry) in stored.iter().enumerate() {
nodes.insert(entry.id, (idx, false));
}
for parent in stored.iter().filter_map(|entry| entry.parent) {
if let Some(node) = nodes.get_mut(&parent) {
node.1 = true;
}
}
let nodes = &nodes;
let is_leaf = |entry: &&Entry| nodes.get(&entry.id).is_some_and(|node| !node.1);
// Root-ward walk; a missing parent ends the recoverable chain, and the
// step cap ends a corrupt cycle.
let chain = |leaf: &Entry| {
let mut cursor = Some(leaf.id);
std::iter::from_fn(move || {
let (idx, _) = nodes.get(&cursor?)?;
let entry = stored.get(*idx)?;
cursor = entry.parent;
Some(*entry)
})
.take(stored.len())
};
let mut branches = Vec::with_capacity(stored.iter().filter(is_leaf).count());
let mut active_branch_idx = 0;
for leaf in stored.iter().filter(is_leaf) {
let mut entries = Vec::with_capacity(chain(leaf).count());
entries.extend(chain(leaf));
entries.reverse();
let is_active = leaf.id == last.id;
if is_active {
active_branch_idx = branches.len();
}
branches.push(Branch {
// The selected branch owns the persisted identity; sibling
// branches derive theirs from their leaf, which keeps every
// reload of the same file identical.
id: if is_active {
active_id
} else {
BranchId::from_u128(leaf.id.as_u128())
},
selected_idx: entries.len().saturating_sub(1),
entries,
});
}
(branches, active_branch_idx)
}
#[cfg(test)]
#[expect(
clippy::unwrap_used,
clippy::panic,
reason = "tests may panic on invariant violations"
)]
mod tests {
use super::*;
fn message(text: &str) -> EntryContent {
EntryContent::Message(Message::with_text(Role::User, text))
}
fn compaction(detail: serde_json::Value) -> EntryContent {
EntryContent::Meta {
kind: COMPACTION_KIND.to_string(),
detail,
}
}
fn text_of(session: &Session, entry: &Entry) -> String {
match session.content(entry).unwrap() {
EntryContent::Message(m) => format!("{}", m.blocks.len()),
EntryContent::Meta { kind, .. } => kind,
other => panic!("unexpected content {other:?}"),
}
}
/// The frame of every record in `entries`, content copied from `session`.
fn frames_of(session: &Session, entries: &[Entry]) -> Vec<Frame> {
let mut frames = vec![Frame::Header {
header: session.header(),
}];
frames.extend(entries.iter().map(|entry| Frame::Known {
entry: EntryFrame {
id: entry.id,
parent: entry.parent,
timestamp_ms: entry.timestamp_ms,
kind: entry.kind,
content: session.bytes(entry).to_vec(),
},
}));
frames
}
fn frames(session: &Session) -> Vec<Frame> {
frames_of(session, &session.active_branch().entries)
}
/// Frames of entries built outside any session, under a fresh header;
/// the only way to store parent links a session would not append.
fn stored(entries: &[EntryFrame]) -> (SessionHeader, Vec<Frame>) {
let header = SessionHeader {
session_id: SessionId::new(),
branch_id: BranchId::new(),
};
let mut frames = vec![Frame::Header { header }];
frames.extend(entries.iter().cloned().map(|entry| Frame::Known { entry }));
(header, frames)
}
fn ids(branch: &Branch) -> Vec<EntryId> {
branch.entries.iter().map(|entry| entry.id).collect()
}
#[test]
fn sibling_entries_restore_as_separate_branches() {
// Stored parent links record two children of one root.
let root = EntryFrame::new(None, &message("root")).unwrap();
let first = EntryFrame::new(Some(root.id), &message("first")).unwrap();
let second = EntryFrame::new(Some(root.id), &message("second")).unwrap();
let (header, stored) = stored(&[root.clone(), first.clone(), second.clone()]);
let session = Session::from_frames(&stored).unwrap();
assert_eq!(session.branches.len(), 2);
// The last stored entry selects the active branch and keeps the
// persisted branch identity.
assert_eq!(session.selected_id(), Some(second.id));
assert_eq!(session.active_branch().id, header.branch_id);
assert_eq!(ids(session.active_branch()), vec![root.id, second.id]);
let sibling = &session.branches[1 - session.active_branch_idx];
assert_eq!(ids(sibling), vec![root.id, first.id]);
// Shared-prefix records point at the same content bytes.
assert_eq!(
session.content(&sibling.entries[0]).unwrap(),
EntryContent::decode(&root.content).unwrap()
);
// Reconstruction is identical on a second pass, sibling ids included.
let again = Session::from_frames(&stored).unwrap();
assert_eq!(again, session);
}
#[test]
fn appended_content_decodes_from_the_arena_and_records_its_kind() {
let mut session = Session::new();
let content = EntryContent::ToolResult {
ok: false,
output: "boom".to_string(),
call_id: ToolCallId::new(),
};
let entry = *session.append(&content).unwrap();
assert_eq!(entry.kind, content.kind());
assert_eq!(session.content(&entry).unwrap(), content);
assert_eq!(session.bytes(&entry).len(), entry.content.len as usize);
// A record from elsewhere reads as empty bytes and fails to decode.
let foreign = Entry {
content: Span {
offset: 1 << 20,
len: 4,
},
..entry
};
assert!(session.bytes(&foreign).is_empty());
assert_eq!(
session.content(&foreign).unwrap_err().code(),
"ENTRY_DECODE"
);
}
#[test]
fn failed_persist_leaves_the_session_unchanged() {
let mut session = Session::new();
session.append(&message("kept")).unwrap();
let before = session.clone();
let err = session
.append_with(
&EntryContent::Meta {
kind: COMPACTION_KIND.to_string(),
detail: serde_json::json!({}),
},
|_, bytes| {
assert!(!bytes.is_empty());
Err(SmithError::Cancelled)
},
)
.unwrap_err();
assert_eq!(err.code(), "CANCELLED");
assert_eq!(session, before);
}
#[test]
fn fork_preserves_ids_order_leaf_and_compaction_boundary() {
let mut session = Session::new();
let root = *session.append(&message("root")).unwrap();
let boundary = *session
.append(&compaction(serde_json::json!({"summarized": 1})))
.unwrap();
let after = *session.append(&message("after")).unwrap();
let forked = session.fork_at(boundary.id).unwrap();
assert_ne!(forked.id, session.id);
assert_eq!(ids(forked.active_branch()), vec![root.id, boundary.id]);
assert_eq!(forked.selected_id(), Some(boundary.id));
assert_eq!(forked.compaction_boundary, Some(boundary.id));
assert_eq!(
forked
.active_branch()
.entries
.iter()
.map(|entry| text_of(&forked, entry))
.collect::<Vec<_>>(),
vec!["1".to_string(), COMPACTION_KIND.to_string()]
);
// The fork carries only the arena prefix its records cover.
assert_eq!(forked.bytes(&after), &[] as &[u8]);
// Forking before the boundary drops it, because it is not retained.
let earlier = session.fork_at(root.id).unwrap();
assert_eq!(earlier.compaction_boundary, None);
assert_eq!(earlier.active_branch().entries.len(), 1);
// The source session is untouched.
assert_eq!(session.active_branch().entries.len(), 3);
assert_eq!(session.selected_id(), Some(after.id));
}
#[test]
fn tool_use_json_inside_messages_encodes_identically() {
let call_id = ToolCallId::new();
let block = |input: serde_json::Value| {
let mut message = Message::with_text(Role::Assistant, "call");
message.add_block(ContentBlock::ToolUse {
name: "write".to_string(),
input,
call_id,
});
message
};
let first = block(serde_json::json!({"path": "p", "content": "c"}));
let mut second = block(serde_json::json!({"content": "c", "path": "p"}));
// Only key order may differ, so encodings must match.
second.id = first.id;
second.timestamp_ms = first.timestamp_ms;
let mut session = Session::new();
let a = *session.append(&EntryContent::Message(first)).unwrap();
let b = *session.append(&EntryContent::Message(second)).unwrap();
assert_eq!(session.bytes(&a), session.bytes(&b));
}
#[test]
fn fork_rejects_unknown_entries() {
let session = Session::new();
let err = session.fork_at(EntryId::new()).unwrap_err();
assert_eq!(err.code(), "ENTRY_NOT_FOUND");
}
#[test]
fn fork_rejects_broken_ancestry_and_forks_sibling_branches() {
// A stored chain whose root parent is missing restores, but cannot fork.
let orphan = EntryFrame::new(Some(EntryId::new()), &message("orphan")).unwrap();
let child = EntryFrame::new(Some(orphan.id), &message("child")).unwrap();
let broken = Session::from_frames(&stored(&[orphan, child.clone()]).1).unwrap();
assert_eq!(broken.selected_id(), Some(child.id));
let err = broken.fork_at(child.id).unwrap_err();
assert_eq!(err.code(), "ENTRY_NOT_FOUND");
// A leaf on an inactive sibling branch forks along its own chain.
let root = EntryFrame::new(None, &message("root")).unwrap();
let first = EntryFrame::new(Some(root.id), &message("first")).unwrap();
let second = EntryFrame::new(Some(root.id), &message("second")).unwrap();
let siblings =
Session::from_frames(&stored(&[root.clone(), first.clone(), second]).1).unwrap();
let forked = siblings.fork_at(first.id).unwrap();
assert_eq!(ids(forked.active_branch()), vec![root.id, first.id]);
assert_eq!(text_of(&forked, &forked.active_branch().entries[1]), "1");
}
#[test]
fn loading_copies_content_bytes_without_decoding_them() {
let mut scratch = Session::new();
let root = *scratch.append(&message("root")).unwrap();
let mut stored = frames(&scratch);
if let Frame::Known { entry } = &mut stored[1] {
entry.content = vec![0xFF, 0x00];
}
let restored = Session::from_frames(&stored).unwrap();
let record = restored.active_branch().entries[0];
assert_eq!(record.kind, root.kind);
assert_eq!(restored.bytes(&record), &[0xFF, 0x00]);
assert_eq!(
restored.content(&record).unwrap_err().code(),
"ENTRY_DECODE"
);
}
#[test]
fn content_of_a_record_whose_kind_lies_is_a_kind_mismatch() {
let mut scratch = Session::new();
scratch.append(&message("root")).unwrap();
let mut stored = frames(&scratch);
let lie = EntryKind::ToolCall {
call_id: ToolCallId::new(),
};
if let Frame::Known { entry } = &mut stored[1] {
entry.kind = lie;
}
// Loading trusts the record; only a decode can tell.
let restored = Session::from_frames(&stored).unwrap();
let record = restored.active_branch().entries[0];
assert_eq!(record.kind, lie);
assert_eq!(
restored.content(&record).unwrap_err().code(),
"ENTRY_KIND_MISMATCH"
);
}
#[test]
fn compaction_boundary_survives_a_frame_round_trip() {
let mut session = Session::new();
session.append(&message("root")).unwrap();
let boundary = *session.append(&compaction(serde_json::json!({}))).unwrap();
let restored = Session::from_frames(&frames(&session)).unwrap();
assert_eq!(restored.compaction_boundary, Some(boundary.id));
assert_eq!(restored.id, session.id);
assert_eq!(restored.active_branch().id, session.header().branch_id);
assert_eq!(restored, session);
}
}
/// Laws of canonicalization, span arithmetic, and content encoding
/// (`SMH-SPEC-SPEC0001`, Sessions).
///
/// `serde_json` preserves insertion order but compares maps order-blind, so
/// key-order laws compare serialized bytes, never values.
#[cfg(test)]
pub mod properties {
use super::*;
use proptest::collection::vec;
use proptest::prelude::*;
use serde_json::Value;
/// Arbitrary JSON, at most three levels deep and 24 nodes large.
fn json() -> impl Strategy<Value = Value> {
let leaf = prop_oneof![
Just(Value::Null),
any::<bool>().prop_map(Value::from),
any::<i64>().prop_map(Value::from),
any::<u64>().prop_map(Value::from),
(proptest::num::f64::NORMAL | proptest::num::f64::ZERO).prop_map(Value::from),
"\\PC{0,8}".prop_map(Value::from),
];
leaf.prop_recursive(3, 24, 4, |inner| {
prop_oneof![
vec(inner.clone(), 0..4).prop_map(Value::Array),
vec(("\\PC{0,4}", inner), 0..5)
.prop_map(|entries| Value::Object(entries.into_iter().collect())),
]
})
}
/// `value` with the keys of every object reordered as `seed` dictates.
fn reordered(value: &Value, seed: u64) -> Value {
match value {
Value::Object(map) => {
let mut entries: Vec<(String, Value)> = map
.iter()
.map(|(key, value)| (key.clone(), reordered(value, seed.rotate_left(7))))
.collect();
if !entries.is_empty() {
let by = usize::try_from(seed % entries.len() as u64).unwrap_or(0);
entries.rotate_left(by);
}
if seed & 1 == 1 {
entries.reverse();
}
Value::Object(entries.into_iter().collect())
}
Value::Array(items) => Value::Array(
items
.iter()
.map(|item| reordered(item, seed.rotate_left(7)))
.collect(),
),
other => other.clone(),
}
}
fn json_bytes(value: &Value) -> std::result::Result<Vec<u8>, TestCaseError> {
serde_json::to_vec(value).map_err(|e| TestCaseError::fail(e.to_string()))
}
fn cbor<T: Serialize>(value: &T) -> std::result::Result<Vec<u8>, TestCaseError> {
let mut out = Vec::new();
ciborium::ser::into_writer(value, &mut out)
.map(|()| out)
.map_err(|e| TestCaseError::fail(e.to_string()))
}
/// Any message role.
pub fn role() -> impl Strategy<Value = Role> {
prop_oneof![
Just(Role::User),
Just(Role::Assistant),
Just(Role::System),
Just(Role::Tool),
]
}
fn block() -> impl Strategy<Value = ContentBlock> {
prop_oneof![
"\\PC{0,16}".prop_map(ContentBlock::Text),
"\\PC{0,16}".prop_map(ContentBlock::Thinking),
("\\PC{0,8}", json(), any::<u128>()).prop_map(|(name, input, id)| {
ContentBlock::tool_use(name, input, ToolCallId::from_u128(id))
}),
(any::<u128>(), any::<bool>(), "\\PC{0,16}").prop_map(|(id, ok, output)| {
ContentBlock::tool_result(ToolCallId::from_u128(id), ok, output)
}),
]
}
fn content() -> impl Strategy<Value = EntryContent> {
prop_oneof![
(any::<u128>(), role(), vec(block(), 0..4), any::<i64>()).prop_map(
|(id, role, blocks, timestamp_ms)| EntryContent::Message(Message {
id: MessageId::from_u128(id),
role,
blocks,
timestamp_ms,
})
),
("\\PC{0,8}", json(), any::<u128>()).prop_map(|(name, input, id)| {
EntryContent::ToolCall {
name,
input,
call_id: ToolCallId::from_u128(id),
}
}),
(any::<bool>(), "\\PC{0,16}", any::<u128>()).prop_map(|(ok, output, id)| {
EntryContent::ToolResult {
ok,
output,
call_id: ToolCallId::from_u128(id),
}
}),
(
prop_oneof![
Just(COMPACTION_KIND.to_string()),
Just(PROVIDER_HTTP_KIND.to_string()),
"\\PC{0,8}",
],
json()
)
.prop_map(|(kind, detail)| EntryContent::Meta { kind, detail }),
]
}
/// `content` with every JSON value in canonical key order.
fn canonical_content(content: &EntryContent) -> EntryContent {
match content {
EntryContent::Message(message) => EntryContent::Message(canonical_message_of(message)),
EntryContent::ToolCall {
name,
input,
call_id,
} => EntryContent::ToolCall {
name: name.clone(),
input: canonical_json(input),
call_id: *call_id,
},
EntryContent::ToolResult { .. } => content.clone(),
EntryContent::Meta { kind, detail } => EntryContent::Meta {
kind: kind.clone(),
detail: canonical_json(detail),
},
}
}
/// Offsets and lengths around the 4 GiB arena cap, plus anything.
fn arena_size() -> impl Strategy<Value = usize> {
let cap = u32::MAX as usize;
prop_oneof![
0..=cap,
cap - 64..=cap + 64,
(0..=64_usize).prop_map(move |below| cap - below),
any::<usize>(),
]
}
proptest! {
#[test]
fn canonical_json_is_idempotent_and_keeps_the_value(value in json()) {
let once = canonical_json(&value);
prop_assert_eq!(json_bytes(&canonical_json(&once))?, json_bytes(&once)?);
prop_assert!(is_canonical(&once));
prop_assert_eq!(once, value);
}
#[test]
fn canonical_json_ignores_key_order(value in json(), seed in any::<u64>()) {
let shuffled = reordered(&value, seed);
prop_assert_eq!(
json_bytes(&canonical_json(&shuffled))?,
json_bytes(&canonical_json(&value))?
);
}
#[test]
fn is_canonical_holds_exactly_when_canonicalizing_changes_nothing(value in json()) {
prop_assert_eq!(
is_canonical(&value),
json_bytes(&canonical_json(&value))? == json_bytes(&value)?
);
}
#[test]
fn span_at_is_exact_below_the_cap_and_an_error_above(
offset in arena_size(),
len in arena_size(),
) {
let fits = offset
.checked_add(len)
.is_some_and(|end| u32::try_from(end).is_ok());
match Span::at(offset, len) {
Ok(span) => {
prop_assert!(fits);
prop_assert_eq!(span.offset as usize, offset);
prop_assert_eq!(span.len as usize, len);
prop_assert_eq!(span.end(), offset + len);
prop_assert_eq!(span.range().len(), len);
}
Err(error) => {
prop_assert!(!fits);
prop_assert_eq!(error.code(), "SESSION_ARENA_FULL");
}
}
}
#[test]
fn content_ref_serializes_like_owned_content(content in content()) {
prop_assert_eq!(cbor(&EntryContentRef::from(&content))?, cbor(&content)?);
}
#[test]
fn owned_and_borrowed_encodings_are_canonical_and_decode_back(content in content()) {
let owned = EntryFrame::new(None, &content)
.map_err(|e| TestCaseError::fail(e.to_string()))?
.content;
let mut session = Session::new();
let record = *session
.append(EntryContentRef::from(&content))
.map_err(|e| TestCaseError::fail(e.to_string()))?;
prop_assert_eq!(session.bytes(&record), owned.as_slice());
prop_assert_eq!(&owned, &cbor(&canonical_content(&content))?);
prop_assert_eq!(
EntryContent::decode(&owned).map_err(|e| TestCaseError::fail(e.to_string()))?,
content
);
}
}
}