session-storeとして分離
This commit is contained in:
@@ -0,0 +1,21 @@
|
||||
//! Debug-only raw stream event recording.
|
||||
//!
|
||||
//! [`TraceEntry`] captures every LLM stream event verbatim for debugging
|
||||
//! and post-hoc analysis. Written to a separate `.trace.jsonl` file,
|
||||
//! completely independent of the session log used for state restoration.
|
||||
//!
|
||||
//! Disabled by default. Enable via `SessionConfig::record_event_trace`.
|
||||
|
||||
use llm_worker::llm_client::event::Event;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// A single trace entry recording a raw stream event.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TraceEntry {
|
||||
/// Timestamp in milliseconds since Unix epoch.
|
||||
pub ts: u64,
|
||||
/// Turn number at the time of recording.
|
||||
pub turn: usize,
|
||||
/// The raw stream event.
|
||||
pub event: Event,
|
||||
}
|
||||
@@ -0,0 +1,157 @@
|
||||
//! Filesystem-backed JSONL store.
|
||||
//!
|
||||
//! Layout:
|
||||
//! - Session log: `{root}/{session_id}.jsonl`
|
||||
//! - Event trace: `{root}/{session_id}.trace.jsonl`
|
||||
|
||||
use crate::event_trace::TraceEntry;
|
||||
use crate::session_log::{EntryHash, HashedEntry};
|
||||
use crate::store::{Store, StoreError};
|
||||
use crate::SessionId;
|
||||
use std::path::{Path, PathBuf};
|
||||
use tokio::fs;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
|
||||
/// Filesystem-backed JSONL store.
|
||||
///
|
||||
/// Each session is stored as a single `.jsonl` file with one [`LogEntry`]
|
||||
/// per line. Writes use append mode for crash safety.
|
||||
#[derive(Clone)]
|
||||
pub struct FsStore {
|
||||
root: PathBuf,
|
||||
}
|
||||
|
||||
impl FsStore {
|
||||
/// Create a new `FsStore` rooted at the given directory.
|
||||
/// Creates the directory if it does not exist.
|
||||
pub async fn new(root: impl Into<PathBuf>) -> Result<Self, StoreError> {
|
||||
let root = root.into();
|
||||
fs::create_dir_all(&root).await?;
|
||||
Ok(Self { root })
|
||||
}
|
||||
|
||||
fn log_path(&self, id: SessionId) -> PathBuf {
|
||||
self.root.join(format!("{id}.jsonl"))
|
||||
}
|
||||
|
||||
fn trace_path(&self, id: SessionId) -> PathBuf {
|
||||
self.root.join(format!("{id}.trace.jsonl"))
|
||||
}
|
||||
|
||||
async fn append_line(&self, path: &Path, line: &str) -> Result<(), StoreError> {
|
||||
let mut file = fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(path)
|
||||
.await?;
|
||||
file.write_all(line.as_bytes()).await?;
|
||||
file.write_all(b"\n").await?;
|
||||
file.flush().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_jsonl<T: serde::de::DeserializeOwned>(
|
||||
content: &str,
|
||||
) -> Result<Vec<T>, StoreError> {
|
||||
let mut entries = Vec::new();
|
||||
for (i, line) in content.lines().enumerate() {
|
||||
if line.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
let entry: T =
|
||||
serde_json::from_str(line).map_err(|e| StoreError::Corrupt {
|
||||
line: i + 1,
|
||||
message: e.to_string(),
|
||||
})?;
|
||||
entries.push(entry);
|
||||
}
|
||||
Ok(entries)
|
||||
}
|
||||
}
|
||||
|
||||
impl Store for FsStore {
|
||||
async fn append(&self, id: SessionId, entry: &HashedEntry) -> Result<(), StoreError> {
|
||||
let line = serde_json::to_string(entry)?;
|
||||
self.append_line(&self.log_path(id), &line).await
|
||||
}
|
||||
|
||||
async fn read_all(&self, id: SessionId) -> Result<Vec<HashedEntry>, StoreError> {
|
||||
let path = self.log_path(id);
|
||||
if !path.exists() {
|
||||
return Err(StoreError::NotFound(id));
|
||||
}
|
||||
let content = fs::read_to_string(&path).await?;
|
||||
Self::parse_jsonl(&content)
|
||||
}
|
||||
|
||||
async fn list_sessions(&self) -> Result<Vec<SessionId>, StoreError> {
|
||||
let mut sessions = Vec::new();
|
||||
let mut dir = fs::read_dir(&self.root).await?;
|
||||
while let Some(entry) = dir.next_entry().await? {
|
||||
let path = entry.path();
|
||||
// Only match .jsonl files, not .trace.jsonl
|
||||
let name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
|
||||
if name.ends_with(".jsonl") && !name.ends_with(".trace.jsonl") {
|
||||
let stem = name.trim_end_matches(".jsonl");
|
||||
if let Ok(id) = stem.parse::<SessionId>() {
|
||||
sessions.push(id);
|
||||
}
|
||||
}
|
||||
}
|
||||
// UUID v7: lexicographic sort = chronological sort, newest first
|
||||
sessions.sort_by(|a, b| b.cmp(a));
|
||||
Ok(sessions)
|
||||
}
|
||||
|
||||
async fn create_session(
|
||||
&self,
|
||||
id: SessionId,
|
||||
entries: &[HashedEntry],
|
||||
) -> Result<(), StoreError> {
|
||||
let path = self.log_path(id);
|
||||
let mut content = String::new();
|
||||
for entry in entries {
|
||||
content.push_str(&serde_json::to_string(entry)?);
|
||||
content.push('\n');
|
||||
}
|
||||
fs::write(&path, content.as_bytes()).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn exists(&self, id: SessionId) -> Result<bool, StoreError> {
|
||||
Ok(self.log_path(id).exists())
|
||||
}
|
||||
|
||||
async fn read_head_hash(
|
||||
&self,
|
||||
id: SessionId,
|
||||
) -> Result<Option<EntryHash>, StoreError> {
|
||||
let path = self.log_path(id);
|
||||
if !path.exists() {
|
||||
return Err(StoreError::NotFound(id));
|
||||
}
|
||||
let content = fs::read_to_string(&path).await?;
|
||||
let last_line = content.lines().rev().find(|l| !l.trim().is_empty());
|
||||
match last_line {
|
||||
Some(line) => {
|
||||
let entry: HashedEntry = serde_json::from_str(line).map_err(|e| {
|
||||
StoreError::Corrupt {
|
||||
line: content.lines().count(),
|
||||
message: e.to_string(),
|
||||
}
|
||||
})?;
|
||||
Ok(Some(entry.hash))
|
||||
}
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn append_trace(
|
||||
&self,
|
||||
id: SessionId,
|
||||
entry: &TraceEntry,
|
||||
) -> Result<(), StoreError> {
|
||||
let line = serde_json::to_string(entry)?;
|
||||
self.append_line(&self.trace_path(id), &line).await
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
//! Session persistence via append-only JSONL logs.
|
||||
//!
|
||||
//! # Architecture
|
||||
//!
|
||||
//! Sessions are recorded as a sequence of [`LogEntry`] values, one per line
|
||||
//! in a `.jsonl` file. Reading the log and collecting entries reconstructs
|
||||
//! the full Worker state — no separate snapshots or checkpoints needed.
|
||||
//!
|
||||
//! This crate provides free functions for persistence operations.
|
||||
//! The caller (typically Pod) holds the Worker directly and calls these
|
||||
//! functions after state-mutating operations.
|
||||
//!
|
||||
//! Debug-mode [`TraceEntry`] records capture raw stream events in a separate
|
||||
//! `.trace.jsonl` file, independent of the session log.
|
||||
//!
|
||||
//! # Quick start
|
||||
//!
|
||||
//! ```ignore
|
||||
//! use session_store::{create_session, restore, save_delta, FsStore, SessionStartState};
|
||||
//!
|
||||
//! let store = FsStore::new("./sessions").await?;
|
||||
//! let (session_id, head_hash) = create_session(&store, SessionStartState {
|
||||
//! system_prompt: None,
|
||||
//! config: &config,
|
||||
//! history: &[],
|
||||
//! }).await?;
|
||||
//! ```
|
||||
|
||||
pub mod event_trace;
|
||||
pub mod fs_store;
|
||||
pub mod session;
|
||||
pub mod session_log;
|
||||
pub mod store;
|
||||
|
||||
pub use event_trace::TraceEntry;
|
||||
pub use fs_store::FsStore;
|
||||
pub use session::{
|
||||
SessionStartState, create_session, ensure_head_or_fork, fork, fork_at, restore, save_cache_locked,
|
||||
save_cache_unlocked, save_config_changed, save_delta, save_outcome, save_turn_end,
|
||||
};
|
||||
pub use session_log::{
|
||||
EntryHash, HashedEntry, LogEntry, Outcome, RestoredState, build_chain, collect_state,
|
||||
compute_hash,
|
||||
};
|
||||
pub use store::{Store, StoreError};
|
||||
|
||||
/// Session identifier. UUID v7 (time-ordered, lexicographically sortable).
|
||||
pub type SessionId = uuid::Uuid;
|
||||
|
||||
/// Generate a new session ID.
|
||||
pub fn new_session_id() -> SessionId {
|
||||
uuid::Uuid::now_v7()
|
||||
}
|
||||
@@ -0,0 +1,291 @@
|
||||
//! Free functions for session persistence operations.
|
||||
//!
|
||||
//! These functions record and restore session state without owning a Worker.
|
||||
//! The caller (typically Pod) holds the Worker directly and calls these
|
||||
//! functions after state-mutating operations.
|
||||
|
||||
use crate::session_log::{self, EntryHash, HashedEntry, LogEntry, Outcome};
|
||||
use crate::store::{Store, StoreError};
|
||||
use crate::SessionId;
|
||||
use llm_worker::llm_client::types::Item;
|
||||
use llm_worker::llm_client::RequestConfig;
|
||||
|
||||
/// State snapshot for creating a SessionStart entry.
|
||||
pub struct SessionStartState<'a> {
|
||||
pub system_prompt: Option<&'a str>,
|
||||
pub config: &'a RequestConfig,
|
||||
pub history: &'a [Item],
|
||||
}
|
||||
|
||||
/// Create a new session, writing the initial `SessionStart` entry.
|
||||
///
|
||||
/// Returns the new session ID and head hash.
|
||||
pub async fn create_session(
|
||||
store: &impl Store,
|
||||
state: SessionStartState<'_>,
|
||||
) -> Result<(SessionId, EntryHash), StoreError> {
|
||||
let session_id = crate::new_session_id();
|
||||
let entry = LogEntry::SessionStart {
|
||||
ts: session_log::now_millis(),
|
||||
system_prompt: state.system_prompt.map(String::from),
|
||||
config: state.config.clone(),
|
||||
history: state.history.to_vec(),
|
||||
};
|
||||
let hash = session_log::compute_hash(None, &entry);
|
||||
let hashed_entry = HashedEntry {
|
||||
hash: hash.clone(),
|
||||
prev_hash: None,
|
||||
entry,
|
||||
};
|
||||
store.append(session_id, &hashed_entry).await?;
|
||||
Ok((session_id, hash))
|
||||
}
|
||||
|
||||
/// Restore session state from a stored log.
|
||||
///
|
||||
/// Returns the reconstructed state. The caller is responsible for
|
||||
/// applying it to a Worker.
|
||||
pub async fn restore(
|
||||
store: &impl Store,
|
||||
session_id: SessionId,
|
||||
) -> Result<crate::session_log::RestoredState, StoreError> {
|
||||
let entries = store.read_all(session_id).await?;
|
||||
Ok(session_log::collect_state(&entries))
|
||||
}
|
||||
|
||||
/// Check if the store's head still matches the expected head hash.
|
||||
/// If not, auto-fork into a new session.
|
||||
///
|
||||
/// Updates `session_id` and `head_hash` in place when a fork occurs.
|
||||
pub async fn ensure_head_or_fork(
|
||||
store: &impl Store,
|
||||
session_id: &mut SessionId,
|
||||
head_hash: &mut Option<EntryHash>,
|
||||
state: SessionStartState<'_>,
|
||||
) -> Result<(), StoreError> {
|
||||
let store_head = store.read_head_hash(*session_id).await?;
|
||||
if store_head == *head_hash {
|
||||
return Ok(());
|
||||
}
|
||||
let fork_id = crate::new_session_id();
|
||||
let entry = LogEntry::SessionStart {
|
||||
ts: session_log::now_millis(),
|
||||
system_prompt: state.system_prompt.map(String::from),
|
||||
config: state.config.clone(),
|
||||
history: state.history.to_vec(),
|
||||
};
|
||||
let hash = session_log::compute_hash(None, &entry);
|
||||
let hashed_entry = HashedEntry {
|
||||
hash: hash.clone(),
|
||||
prev_hash: None,
|
||||
entry,
|
||||
};
|
||||
store.create_session(fork_id, &[hashed_entry]).await?;
|
||||
*session_id = fork_id;
|
||||
*head_hash = Some(hash);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Log the history delta — new items added since the previous snapshot.
|
||||
///
|
||||
/// Classifies items into UserInput, AssistantItems, ToolResults, and
|
||||
/// HookInjectedItems entries automatically.
|
||||
pub async fn save_delta(
|
||||
store: &impl Store,
|
||||
session_id: SessionId,
|
||||
head_hash: &mut Option<EntryHash>,
|
||||
new_items: &[Item],
|
||||
) -> Result<(), StoreError> {
|
||||
if new_items.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let ts = session_log::now_millis();
|
||||
let mut i = 0;
|
||||
|
||||
while i < new_items.len() {
|
||||
let item = &new_items[i];
|
||||
if item.is_user_message() {
|
||||
append_entry(store, session_id, head_hash, LogEntry::UserInput {
|
||||
ts,
|
||||
item: new_items[i].clone(),
|
||||
})
|
||||
.await?;
|
||||
i += 1;
|
||||
} else if item.is_tool_result() {
|
||||
let start = i;
|
||||
while i < new_items.len() && new_items[i].is_tool_result() {
|
||||
i += 1;
|
||||
}
|
||||
append_entry(store, session_id, head_hash, LogEntry::ToolResults {
|
||||
ts,
|
||||
items: new_items[start..i].to_vec(),
|
||||
})
|
||||
.await?;
|
||||
} else if item.is_assistant_message() || item.is_tool_call() || item.is_reasoning() {
|
||||
let start = i;
|
||||
while i < new_items.len()
|
||||
&& (new_items[i].is_assistant_message()
|
||||
|| new_items[i].is_tool_call()
|
||||
|| new_items[i].is_reasoning())
|
||||
{
|
||||
i += 1;
|
||||
}
|
||||
append_entry(store, session_id, head_hash, LogEntry::AssistantItems {
|
||||
ts,
|
||||
items: new_items[start..i].to_vec(),
|
||||
})
|
||||
.await?;
|
||||
} else {
|
||||
append_entry(store, session_id, head_hash, LogEntry::HookInjectedItems {
|
||||
ts,
|
||||
items: vec![new_items[i].clone()],
|
||||
})
|
||||
.await?;
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Log a TurnEnd entry.
|
||||
pub async fn save_turn_end(
|
||||
store: &impl Store,
|
||||
session_id: SessionId,
|
||||
head_hash: &mut Option<EntryHash>,
|
||||
turn_count: usize,
|
||||
) -> Result<(), StoreError> {
|
||||
append_entry(store, session_id, head_hash, LogEntry::TurnEnd {
|
||||
ts: session_log::now_millis(),
|
||||
turn_count,
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Log a RunOutcome entry.
|
||||
pub async fn save_outcome(
|
||||
store: &impl Store,
|
||||
session_id: SessionId,
|
||||
head_hash: &mut Option<EntryHash>,
|
||||
outcome: Outcome,
|
||||
interrupted: bool,
|
||||
) -> Result<(), StoreError> {
|
||||
append_entry(store, session_id, head_hash, LogEntry::RunOutcome {
|
||||
ts: session_log::now_millis(),
|
||||
outcome,
|
||||
interrupted,
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Log a `Locked` entry (KV cache locked).
|
||||
pub async fn save_cache_locked(
|
||||
store: &impl Store,
|
||||
session_id: SessionId,
|
||||
head_hash: &mut Option<EntryHash>,
|
||||
locked_prefix_len: usize,
|
||||
) -> Result<(), StoreError> {
|
||||
append_entry(store, session_id, head_hash, LogEntry::Locked {
|
||||
ts: session_log::now_millis(),
|
||||
locked_prefix_len,
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Log a `CacheUnlocked` entry.
|
||||
pub async fn save_cache_unlocked(
|
||||
store: &impl Store,
|
||||
session_id: SessionId,
|
||||
head_hash: &mut Option<EntryHash>,
|
||||
) -> Result<(), StoreError> {
|
||||
append_entry(store, session_id, head_hash, LogEntry::CacheUnlocked {
|
||||
ts: session_log::now_millis(),
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Log a `ConfigChanged` entry.
|
||||
pub async fn save_config_changed(
|
||||
store: &impl Store,
|
||||
session_id: SessionId,
|
||||
head_hash: &mut Option<EntryHash>,
|
||||
config: &RequestConfig,
|
||||
) -> Result<(), StoreError> {
|
||||
append_entry(store, session_id, head_hash, LogEntry::ConfigChanged {
|
||||
ts: session_log::now_millis(),
|
||||
config: config.clone(),
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Fork the current state into a new session.
|
||||
pub async fn fork(
|
||||
store: &impl Store,
|
||||
state: SessionStartState<'_>,
|
||||
) -> Result<SessionId, StoreError> {
|
||||
let fork_id = crate::new_session_id();
|
||||
let entry = LogEntry::SessionStart {
|
||||
ts: session_log::now_millis(),
|
||||
system_prompt: state.system_prompt.map(String::from),
|
||||
config: state.config.clone(),
|
||||
history: state.history.to_vec(),
|
||||
};
|
||||
let hash = session_log::compute_hash(None, &entry);
|
||||
let hashed_entry = HashedEntry {
|
||||
hash,
|
||||
prev_hash: None,
|
||||
entry,
|
||||
};
|
||||
store.create_session(fork_id, &[hashed_entry]).await?;
|
||||
Ok(fork_id)
|
||||
}
|
||||
|
||||
/// Fork from an arbitrary point in a stored session's log.
|
||||
pub async fn fork_at(
|
||||
store: &impl Store,
|
||||
source_id: SessionId,
|
||||
at_hash: &EntryHash,
|
||||
) -> Result<SessionId, StoreError> {
|
||||
let entries = store.read_all(source_id).await?;
|
||||
let cut = entries
|
||||
.iter()
|
||||
.position(|e| &e.hash == at_hash)
|
||||
.map(|i| i + 1)
|
||||
.unwrap_or(entries.len());
|
||||
let state = session_log::collect_state(&entries[..cut]);
|
||||
|
||||
let fork_id = crate::new_session_id();
|
||||
let entry = LogEntry::SessionStart {
|
||||
ts: session_log::now_millis(),
|
||||
system_prompt: state.system_prompt,
|
||||
config: state.config,
|
||||
history: state.history,
|
||||
};
|
||||
let hash = session_log::compute_hash(None, &entry);
|
||||
let hashed_entry = HashedEntry {
|
||||
hash,
|
||||
prev_hash: None,
|
||||
entry,
|
||||
};
|
||||
store.create_session(fork_id, &[hashed_entry]).await?;
|
||||
Ok(fork_id)
|
||||
}
|
||||
|
||||
// ── Private helper ──────────────────────────────────────────────────────
|
||||
|
||||
async fn append_entry(
|
||||
store: &impl Store,
|
||||
session_id: SessionId,
|
||||
head_hash: &mut Option<EntryHash>,
|
||||
entry: LogEntry,
|
||||
) -> Result<(), StoreError> {
|
||||
let hash = session_log::compute_hash(head_hash.as_ref(), &entry);
|
||||
let hashed_entry = HashedEntry {
|
||||
hash: hash.clone(),
|
||||
prev_hash: head_hash.clone(),
|
||||
entry,
|
||||
};
|
||||
store.append(session_id, &hashed_entry).await?;
|
||||
*head_hash = Some(hash);
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,438 @@
|
||||
//! Session log types for append-only JSONL persistence.
|
||||
//!
|
||||
//! Each [`LogEntry`] represents a single state transition in a session,
|
||||
//! serialized as one line in a `.jsonl` file. Reading all entries and
|
||||
//! collecting them via [`collect_state`] reconstructs the full [`Worker`] state.
|
||||
//!
|
||||
//! Entries are chained via [`EntryHash`]: each [`HashedEntry`] records the hash
|
||||
//! of the previous entry, forming a tamper-evident append-only chain. This
|
||||
//! enables safe fork detection when multiple writers share a session.
|
||||
|
||||
use llm_worker::llm_client::types::{Item, RequestConfig};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
/// SHA-256 hash identifying a specific log entry in the chain.
|
||||
///
|
||||
/// Computed as `sha256(prev_hash_bytes || canonical_json(entry))`.
|
||||
/// Displayed and serialized as a lowercase hex string.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
pub struct EntryHash([u8; 32]);
|
||||
|
||||
impl EntryHash {
|
||||
pub fn as_bytes(&self) -> &[u8; 32] {
|
||||
&self.0
|
||||
}
|
||||
|
||||
pub fn to_hex(&self) -> String {
|
||||
hex::encode(self.0)
|
||||
}
|
||||
|
||||
pub fn from_hex(s: &str) -> Result<Self, hex::FromHexError> {
|
||||
let mut buf = [0u8; 32];
|
||||
hex::decode_to_slice(s, &mut buf)?;
|
||||
Ok(Self(buf))
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for EntryHash {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str(&self.to_hex())
|
||||
}
|
||||
}
|
||||
|
||||
impl Serialize for EntryHash {
|
||||
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
|
||||
serializer.serialize_str(&self.to_hex())
|
||||
}
|
||||
}
|
||||
|
||||
impl<'de> Deserialize<'de> for EntryHash {
|
||||
fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
|
||||
let s = String::deserialize(deserializer)?;
|
||||
Self::from_hex(&s).map_err(serde::de::Error::custom)
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute the hash for a log entry given its predecessor's hash.
|
||||
pub fn compute_hash(prev: Option<&EntryHash>, entry: &LogEntry) -> EntryHash {
|
||||
let mut hasher = Sha256::new();
|
||||
|
||||
// Feed prev_hash bytes (32 zero bytes if None).
|
||||
match prev {
|
||||
Some(h) => hasher.update(h.as_bytes()),
|
||||
None => hasher.update([0u8; 32]),
|
||||
}
|
||||
|
||||
// Canonical JSON of the entry.
|
||||
let json = serde_json::to_string(entry).expect("LogEntry serialization cannot fail");
|
||||
hasher.update(json.as_bytes());
|
||||
|
||||
EntryHash(hasher.finalize().into())
|
||||
}
|
||||
|
||||
/// A [`LogEntry`] with hash-chain metadata.
|
||||
///
|
||||
/// This is the unit persisted to JSONL — one line per `HashedEntry`.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct HashedEntry {
|
||||
pub hash: EntryHash,
|
||||
pub prev_hash: Option<EntryHash>,
|
||||
#[serde(flatten)]
|
||||
pub entry: LogEntry,
|
||||
}
|
||||
|
||||
/// A single session log entry, serialized as one JSONL line.
|
||||
///
|
||||
/// Variants correspond to specific mutation points in `Worker`:
|
||||
/// - `SessionStart` — always the first entry; captures initial state
|
||||
/// - `UserInput` / `AssistantItems` / `ToolResults` / `HookInjectedItems` — history appends
|
||||
/// - `TurnEnd` — turn boundary marker
|
||||
/// - `Locked` / `CacheUnlocked` — KV cache state transitions
|
||||
/// - `RunOutcome` — marks end of a `run()` or `resume()` call
|
||||
/// - `ConfigChanged` — `RequestConfig` mutation
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||
pub enum LogEntry {
|
||||
/// Session start. Always the first entry in a log.
|
||||
/// For forked sessions, `history` contains the seed state from the parent.
|
||||
SessionStart {
|
||||
ts: u64,
|
||||
system_prompt: Option<String>,
|
||||
config: RequestConfig,
|
||||
history: Vec<Item>,
|
||||
},
|
||||
|
||||
/// User input pushed to history (worker.rs:229).
|
||||
UserInput { ts: u64, item: Item },
|
||||
|
||||
/// Assistant response items added to history (worker.rs:1040-1041).
|
||||
AssistantItems { ts: u64, items: Vec<Item> },
|
||||
|
||||
/// Tool execution results added to history (worker.rs:897-900, 1072-1076).
|
||||
ToolResults { ts: u64, items: Vec<Item> },
|
||||
|
||||
/// Items injected by `on_turn_end` hook via `ContinueWithMessages` (worker.rs:1055).
|
||||
HookInjectedItems { ts: u64, items: Vec<Item> },
|
||||
|
||||
/// Turn boundary. Records the turn count after increment.
|
||||
TurnEnd { ts: u64, turn_count: usize },
|
||||
|
||||
/// KV cache locked. Records the history prefix length that is now immutable.
|
||||
#[serde(alias = "cache_locked")]
|
||||
Locked { ts: u64, locked_prefix_len: usize },
|
||||
|
||||
/// KV cache unlocked.
|
||||
CacheUnlocked { ts: u64 },
|
||||
|
||||
/// Outcome of a `run()` or `resume()` call.
|
||||
/// This is metadata for auditing; state collection does not branch on the outcome.
|
||||
RunOutcome {
|
||||
ts: u64,
|
||||
outcome: Outcome,
|
||||
interrupted: bool,
|
||||
},
|
||||
|
||||
/// `RequestConfig` changed.
|
||||
ConfigChanged { ts: u64, config: RequestConfig },
|
||||
}
|
||||
|
||||
/// Outcome of a run/resume call. Metadata for auditing only.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum Outcome {
|
||||
Finished,
|
||||
Paused,
|
||||
LimitReached,
|
||||
Error { message: String },
|
||||
}
|
||||
|
||||
/// State collected from log entries.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RestoredState {
|
||||
pub system_prompt: Option<String>,
|
||||
pub config: RequestConfig,
|
||||
pub history: Vec<Item>,
|
||||
pub turn_count: usize,
|
||||
pub locked_prefix_len: usize,
|
||||
pub last_run_interrupted: bool,
|
||||
/// Hash of the last entry in the chain (None if empty).
|
||||
pub head_hash: Option<EntryHash>,
|
||||
}
|
||||
|
||||
/// Replay a sequence of hashed entries to reconstruct worker state.
|
||||
pub fn collect_state(entries: &[HashedEntry]) -> RestoredState {
|
||||
let mut state = RestoredState {
|
||||
system_prompt: None,
|
||||
config: RequestConfig::default(),
|
||||
history: Vec::new(),
|
||||
turn_count: 0,
|
||||
locked_prefix_len: 0,
|
||||
last_run_interrupted: false,
|
||||
head_hash: None,
|
||||
};
|
||||
|
||||
for hashed in entries {
|
||||
state.head_hash = Some(hashed.hash.clone());
|
||||
|
||||
match &hashed.entry {
|
||||
LogEntry::SessionStart {
|
||||
system_prompt,
|
||||
config,
|
||||
history,
|
||||
..
|
||||
} => {
|
||||
state.system_prompt = system_prompt.clone();
|
||||
state.config = config.clone();
|
||||
state.history = history.clone();
|
||||
}
|
||||
LogEntry::UserInput { item, .. } => {
|
||||
state.history.push(item.clone());
|
||||
}
|
||||
LogEntry::AssistantItems { items, .. } => {
|
||||
state.history.extend(items.iter().cloned());
|
||||
}
|
||||
LogEntry::ToolResults { items, .. } => {
|
||||
state.history.extend(items.iter().cloned());
|
||||
}
|
||||
LogEntry::HookInjectedItems { items, .. } => {
|
||||
state.history.extend(items.iter().cloned());
|
||||
}
|
||||
LogEntry::TurnEnd { turn_count, .. } => {
|
||||
state.turn_count = *turn_count;
|
||||
}
|
||||
LogEntry::Locked {
|
||||
locked_prefix_len, ..
|
||||
} => {
|
||||
state.locked_prefix_len = *locked_prefix_len;
|
||||
}
|
||||
LogEntry::CacheUnlocked { .. } => {
|
||||
state.locked_prefix_len = 0;
|
||||
}
|
||||
LogEntry::RunOutcome { interrupted, .. } => {
|
||||
state.last_run_interrupted = *interrupted;
|
||||
}
|
||||
LogEntry::ConfigChanged { config, .. } => {
|
||||
state.config = config.clone();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
state
|
||||
}
|
||||
|
||||
/// Get the current timestamp in milliseconds since Unix epoch.
|
||||
pub fn now_millis() -> u64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.expect("system clock before Unix epoch")
|
||||
.as_millis() as u64
|
||||
}
|
||||
|
||||
/// Build a hash chain from plain `LogEntry` values.
|
||||
///
|
||||
/// Useful for tests and for seeding new sessions from a list of entries.
|
||||
pub fn build_chain(entries: &[LogEntry]) -> Vec<HashedEntry> {
|
||||
let mut chain = Vec::with_capacity(entries.len());
|
||||
let mut prev: Option<EntryHash> = None;
|
||||
|
||||
for entry in entries {
|
||||
let hash = compute_hash(prev.as_ref(), entry);
|
||||
chain.push(HashedEntry {
|
||||
hash: hash.clone(),
|
||||
prev_hash: prev,
|
||||
entry: entry.clone(),
|
||||
});
|
||||
prev = Some(hash);
|
||||
}
|
||||
|
||||
chain
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn replay_empty() {
|
||||
let state = collect_state(&[]);
|
||||
assert!(state.history.is_empty());
|
||||
assert_eq!(state.turn_count, 0);
|
||||
assert_eq!(state.locked_prefix_len, 0);
|
||||
assert!(state.head_hash.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replay_session_start_sets_initial_state() {
|
||||
let entries = build_chain(&[LogEntry::SessionStart {
|
||||
ts: 1000,
|
||||
system_prompt: Some("You are helpful.".into()),
|
||||
config: RequestConfig::default().with_max_tokens(1024),
|
||||
history: vec![Item::user_message("seed")],
|
||||
}]);
|
||||
let state = collect_state(&entries);
|
||||
assert_eq!(state.system_prompt.as_deref(), Some("You are helpful."));
|
||||
assert_eq!(state.config.max_tokens, Some(1024));
|
||||
assert_eq!(state.history.len(), 1);
|
||||
assert!(state.head_hash.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replay_full_turn() {
|
||||
let entries = build_chain(&[
|
||||
LogEntry::SessionStart {
|
||||
ts: 1000,
|
||||
system_prompt: None,
|
||||
config: RequestConfig::default(),
|
||||
history: vec![],
|
||||
},
|
||||
LogEntry::UserInput {
|
||||
ts: 2000,
|
||||
item: Item::user_message("Hello"),
|
||||
},
|
||||
LogEntry::AssistantItems {
|
||||
ts: 3000,
|
||||
items: vec![Item::assistant_message("Hi!")],
|
||||
},
|
||||
LogEntry::TurnEnd {
|
||||
ts: 3100,
|
||||
turn_count: 1,
|
||||
},
|
||||
LogEntry::RunOutcome {
|
||||
ts: 3200,
|
||||
outcome: Outcome::Finished,
|
||||
interrupted: false,
|
||||
},
|
||||
]);
|
||||
let state = collect_state(&entries);
|
||||
assert_eq!(state.history.len(), 2);
|
||||
assert_eq!(state.turn_count, 1);
|
||||
assert!(!state.last_run_interrupted);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replay_with_tool_calls() {
|
||||
let entries = build_chain(&[
|
||||
LogEntry::SessionStart {
|
||||
ts: 1000,
|
||||
system_prompt: None,
|
||||
config: RequestConfig::default(),
|
||||
history: vec![],
|
||||
},
|
||||
LogEntry::UserInput {
|
||||
ts: 2000,
|
||||
item: Item::user_message("Check weather"),
|
||||
},
|
||||
LogEntry::AssistantItems {
|
||||
ts: 3000,
|
||||
items: vec![Item::tool_call("call_1", "get_weather", r#"{"city":"Tokyo"}"#)],
|
||||
},
|
||||
LogEntry::ToolResults {
|
||||
ts: 3500,
|
||||
items: vec![Item::tool_result("call_1", "Sunny, 25C")],
|
||||
},
|
||||
LogEntry::AssistantItems {
|
||||
ts: 4000,
|
||||
items: vec![Item::assistant_message("It's sunny in Tokyo!")],
|
||||
},
|
||||
LogEntry::TurnEnd {
|
||||
ts: 4100,
|
||||
turn_count: 1,
|
||||
},
|
||||
]);
|
||||
let state = collect_state(&entries);
|
||||
assert_eq!(state.history.len(), 4);
|
||||
assert!(state.history[1].is_tool_call());
|
||||
assert!(state.history[2].is_tool_result());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replay_cache_lock_unlock() {
|
||||
let entries = build_chain(&[
|
||||
LogEntry::SessionStart {
|
||||
ts: 1000,
|
||||
system_prompt: None,
|
||||
config: RequestConfig::default(),
|
||||
history: vec![Item::user_message("a"), Item::assistant_message("b")],
|
||||
},
|
||||
LogEntry::Locked {
|
||||
ts: 2000,
|
||||
locked_prefix_len: 2,
|
||||
},
|
||||
LogEntry::CacheUnlocked { ts: 3000 },
|
||||
]);
|
||||
let state = collect_state(&entries);
|
||||
assert_eq!(state.locked_prefix_len, 0);
|
||||
|
||||
// Check locked state before unlock
|
||||
let state_locked = collect_state(&entries[..2]);
|
||||
assert_eq!(state_locked.locked_prefix_len, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replay_config_changed() {
|
||||
let entries = build_chain(&[
|
||||
LogEntry::SessionStart {
|
||||
ts: 1000,
|
||||
system_prompt: None,
|
||||
config: RequestConfig::default(),
|
||||
history: vec![],
|
||||
},
|
||||
LogEntry::ConfigChanged {
|
||||
ts: 2000,
|
||||
config: RequestConfig::default().with_temperature(0.5),
|
||||
},
|
||||
]);
|
||||
let state = collect_state(&entries);
|
||||
assert_eq!(state.config.temperature, Some(0.5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hash_chain_is_deterministic() {
|
||||
let raw = vec![
|
||||
LogEntry::SessionStart {
|
||||
ts: 1000,
|
||||
system_prompt: None,
|
||||
config: RequestConfig::default(),
|
||||
history: vec![],
|
||||
},
|
||||
LogEntry::UserInput {
|
||||
ts: 2000,
|
||||
item: Item::user_message("Hello"),
|
||||
},
|
||||
];
|
||||
let chain_a = build_chain(&raw);
|
||||
let chain_b = build_chain(&raw);
|
||||
assert_eq!(chain_a[0].hash, chain_b[0].hash);
|
||||
assert_eq!(chain_a[1].hash, chain_b[1].hash);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn different_content_produces_different_hash() {
|
||||
let entry_a = LogEntry::UserInput {
|
||||
ts: 1000,
|
||||
item: Item::user_message("Hello"),
|
||||
};
|
||||
let entry_b = LogEntry::UserInput {
|
||||
ts: 1000,
|
||||
item: Item::user_message("World"),
|
||||
};
|
||||
let hash_a = compute_hash(None, &entry_a);
|
||||
let hash_b = compute_hash(None, &entry_b);
|
||||
assert_ne!(hash_a, hash_b);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hash_hex_round_trip() {
|
||||
let entry = LogEntry::SessionStart {
|
||||
ts: 1000,
|
||||
system_prompt: None,
|
||||
config: RequestConfig::default(),
|
||||
history: vec![],
|
||||
};
|
||||
let hash = compute_hash(None, &entry);
|
||||
let hex = hash.to_hex();
|
||||
let parsed = EntryHash::from_hex(&hex).unwrap();
|
||||
assert_eq!(hash, parsed);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
//! Persistence backend abstraction.
|
||||
//!
|
||||
//! [`Store`] defines the async interface for reading and writing session logs.
|
||||
//! Implementations handle the physical storage (filesystem, database, etc.).
|
||||
|
||||
use crate::event_trace::TraceEntry;
|
||||
use crate::session_log::{EntryHash, HashedEntry};
|
||||
use crate::SessionId;
|
||||
use std::future::Future;
|
||||
|
||||
/// Errors from the persistence store.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum StoreError {
|
||||
#[error("I/O error: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
|
||||
#[error("serialization error: {0}")]
|
||||
Serde(#[from] serde_json::Error),
|
||||
|
||||
#[error("session not found: {0}")]
|
||||
NotFound(SessionId),
|
||||
|
||||
#[error("log corrupted at line {line}: {message}")]
|
||||
Corrupt { line: usize, message: String },
|
||||
}
|
||||
|
||||
/// Async persistence backend for session logs.
|
||||
///
|
||||
/// All methods take `&self` — implementations should use interior mutability
|
||||
/// (e.g., append-mode file handles) when needed.
|
||||
pub trait Store: Send + Sync {
|
||||
/// Append a single hashed entry to the session log.
|
||||
fn append(
|
||||
&self,
|
||||
id: SessionId,
|
||||
entry: &HashedEntry,
|
||||
) -> impl Future<Output = Result<(), StoreError>> + Send;
|
||||
|
||||
/// Read all hashed entries for a session, in order.
|
||||
fn read_all(
|
||||
&self,
|
||||
id: SessionId,
|
||||
) -> impl Future<Output = Result<Vec<HashedEntry>, StoreError>> + Send;
|
||||
|
||||
/// List all session IDs, most recent first.
|
||||
fn list_sessions(
|
||||
&self,
|
||||
) -> impl Future<Output = Result<Vec<SessionId>, StoreError>> + Send;
|
||||
|
||||
/// Create a new session with initial entries.
|
||||
fn create_session(
|
||||
&self,
|
||||
id: SessionId,
|
||||
entries: &[HashedEntry],
|
||||
) -> impl Future<Output = Result<(), StoreError>> + Send;
|
||||
|
||||
/// Check if a session exists.
|
||||
fn exists(
|
||||
&self,
|
||||
id: SessionId,
|
||||
) -> impl Future<Output = Result<bool, StoreError>> + Send;
|
||||
|
||||
/// Read the hash of the last entry in a session (the head).
|
||||
///
|
||||
/// Returns `None` if the session is empty.
|
||||
fn read_head_hash(
|
||||
&self,
|
||||
id: SessionId,
|
||||
) -> impl Future<Output = Result<Option<EntryHash>, StoreError>> + Send;
|
||||
|
||||
/// Append a trace entry to the debug event trace file.
|
||||
fn append_trace(
|
||||
&self,
|
||||
id: SessionId,
|
||||
entry: &TraceEntry,
|
||||
) -> impl Future<Output = Result<(), StoreError>> + Send;
|
||||
}
|
||||
Reference in New Issue
Block a user