session-storeとして分離

This commit is contained in:
2026-04-12 06:31:34 +09:00
parent eb670bfba5
commit cdafd5d914
25 changed files with 910 additions and 657 deletions
+23
View File
@@ -0,0 +1,23 @@
[package]
name = "session-store"
description = "Session persistence via append-only JSONL logs"
version = "0.1.0"
edition.workspace = true
license.workspace = true
[dependencies]
llm-worker = { path = "../llm-worker" }
async-trait = "0.1"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
tokio = { version = "1.49", features = ["fs", "io-util"] }
uuid = { version = "1", features = ["v7", "serde"] }
thiserror = "2.0"
sha2 = "0.11.0"
hex = "0.4.3"
[dev-dependencies]
tokio = { version = "1.49", features = ["macros", "rt-multi-thread", "fs", "io-util"] }
tempfile = "3.24"
futures = "0.3"
async-trait = "0.1"
+29
View File
@@ -0,0 +1,29 @@
# llm-worker-persistence
Worker のセッション永続化を提供するクレート。追記専用の JSONL ログとして状態遷移を記録し、ログの再生によってセッションを完全に復元する。大きなツール出力は Blob ストアに分離保存する。
## 公開型
### セッション
- `Session<C, St>` — Worker をラップした永続化セッション(`run()`, `resume()`, `fork()`, `fork_at()`
- `SessionId` — UUID v7 によるセッション識別子
- `SessionConfig` — 永続化設定(イベントトレース記録の有無)
### ストア
- `Store` トレイト — 永続化バックエンド抽象(`append`, `read_all`, `list_sessions`
- `FsStore` — ファイルシステム上の JSONL ストア実装
- `BlobStore` トレイト — Blob ストレージ抽象(`store`, `load`
- `FsBlobStore` — ファイルシステム上の Blob ストア実装
- `BlobOutputProcessor` — ToolOutputProcessor 実装(小さい出力はインライン、大きい出力は Blob 保存)
### ログ
- `LogEntry` — セッションログのエントリ型(`SessionStart`, `UserInput`, `AssistantItems`, `TurnEnd` など)
- `RestoredState` — ログ再生で復元された状態
- `collect_state()` — ログエントリ列から状態を復元する関数
### ツール
- `InspectTool` — Blob 内容を取得する組み込みツール(行範囲・配列スライス・キー指定セレクタ対応)
+21
View File
@@ -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,
}
+157
View File
@@ -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
}
}
+53
View File
@@ -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()
}
+291
View File
@@ -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(())
}
+438
View File
@@ -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);
}
}
+77
View File
@@ -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;
}
+50
View File
@@ -0,0 +1,50 @@
#![allow(dead_code)]
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use async_trait::async_trait;
use futures::Stream;
use llm_worker::llm_client::event::Event;
use llm_worker::llm_client::{ClientError, LlmClient, Request};
/// A mock LLM client that replays pre-defined event sequences.
#[derive(Clone)]
pub struct MockLlmClient {
responses: Arc<Vec<Vec<Event>>>,
call_count: Arc<AtomicUsize>,
}
impl MockLlmClient {
pub fn new(events: Vec<Event>) -> Self {
Self::with_responses(vec![events])
}
pub fn with_responses(responses: Vec<Vec<Event>>) -> Self {
Self {
responses: Arc::new(responses),
call_count: Arc::new(AtomicUsize::new(0)),
}
}
}
#[async_trait]
impl LlmClient for MockLlmClient {
async fn stream(
&self,
_request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>, ClientError> {
let count = self.call_count.fetch_add(1, Ordering::SeqCst);
if count >= self.responses.len() {
return Err(ClientError::Api {
status: Some(500),
code: Some("mock_error".to_string()),
message: "No more mock responses".to_string(),
});
}
let events = self.responses[count].clone();
let stream = futures::stream::iter(events.into_iter().map(Ok));
Ok(Box::pin(stream))
}
}
+205
View File
@@ -0,0 +1,205 @@
use llm_worker::llm_client::types::{Item, RequestConfig};
use session_store::{
FsStore, LogEntry, Outcome, Store, TraceEntry, build_chain, collect_state, new_session_id,
};
#[tokio::test]
async fn round_trip_write_and_read() {
let dir = tempfile::tempdir().unwrap();
let store = FsStore::new(dir.path()).await.unwrap();
let id = new_session_id();
let raw = vec![
LogEntry::SessionStart {
ts: 1000,
system_prompt: Some("You are helpful.".into()),
config: RequestConfig::default().with_max_tokens(1024),
history: vec![],
},
LogEntry::UserInput {
ts: 2000,
item: Item::user_message("Hello"),
},
LogEntry::AssistantItems {
ts: 3000,
items: vec![Item::assistant_message("Hi there!")],
},
LogEntry::TurnEnd {
ts: 3100,
turn_count: 1,
},
LogEntry::RunOutcome {
ts: 3200,
outcome: Outcome::Finished,
interrupted: false,
},
];
let entries = build_chain(&raw);
// Write entries one by one
for entry in &entries {
store.append(id, entry).await.unwrap();
}
// Read back
let read_back = store.read_all(id).await.unwrap();
assert_eq!(read_back.len(), entries.len());
// Verify hashes survived round-trip
for (orig, read) in entries.iter().zip(read_back.iter()) {
assert_eq!(orig.hash, read.hash);
assert_eq!(orig.prev_hash, read.prev_hash);
}
// Replay and verify state
let state = collect_state(&read_back);
assert_eq!(state.system_prompt.as_deref(), Some("You are helpful."));
assert_eq!(state.config.max_tokens, Some(1024));
assert_eq!(state.history.len(), 2);
assert_eq!(state.turn_count, 1);
assert!(!state.last_run_interrupted);
assert!(state.head_hash.is_some());
}
#[tokio::test]
async fn create_session_writes_all_entries() {
let dir = tempfile::tempdir().unwrap();
let store = FsStore::new(dir.path()).await.unwrap();
let id = new_session_id();
let entries = build_chain(&[LogEntry::SessionStart {
ts: 1000,
system_prompt: None,
config: RequestConfig::default(),
history: vec![Item::user_message("seed"), Item::assistant_message("ok")],
}]);
store.create_session(id, &entries).await.unwrap();
let read_back = store.read_all(id).await.unwrap();
assert_eq!(read_back.len(), 1);
let state = collect_state(&read_back);
assert_eq!(state.history.len(), 2);
}
#[tokio::test]
async fn list_sessions_returns_newest_first() {
let dir = tempfile::tempdir().unwrap();
let store = FsStore::new(dir.path()).await.unwrap();
let id1 = new_session_id();
// Small delay to ensure different UUID v7 timestamps
tokio::time::sleep(std::time::Duration::from_millis(2)).await;
let id2 = new_session_id();
let entries1 = build_chain(&[LogEntry::SessionStart {
ts: 1000,
system_prompt: None,
config: RequestConfig::default(),
history: vec![],
}]);
let entries2 = build_chain(&[LogEntry::SessionStart {
ts: 1001,
system_prompt: None,
config: RequestConfig::default(),
history: vec![],
}]);
store.append(id1, &entries1[0]).await.unwrap();
store.append(id2, &entries2[0]).await.unwrap();
let sessions = store.list_sessions().await.unwrap();
assert_eq!(sessions.len(), 2);
assert_eq!(sessions[0], id2); // newest first
assert_eq!(sessions[1], id1);
}
#[tokio::test]
async fn exists_returns_correct_state() {
let dir = tempfile::tempdir().unwrap();
let store = FsStore::new(dir.path()).await.unwrap();
let id = new_session_id();
assert!(!store.exists(id).await.unwrap());
let entries = build_chain(&[LogEntry::SessionStart {
ts: 1000,
system_prompt: None,
config: RequestConfig::default(),
history: vec![],
}]);
store.append(id, &entries[0]).await.unwrap();
assert!(store.exists(id).await.unwrap());
}
#[tokio::test]
async fn not_found_error_for_missing_session() {
let dir = tempfile::tempdir().unwrap();
let store = FsStore::new(dir.path()).await.unwrap();
let id = new_session_id();
let result = store.read_all(id).await;
assert!(result.is_err());
}
#[tokio::test]
async fn trace_entries_in_separate_file() {
let dir = tempfile::tempdir().unwrap();
let store = FsStore::new(dir.path()).await.unwrap();
let id = new_session_id();
// Write a log entry
let entries = build_chain(&[LogEntry::SessionStart {
ts: 1000,
system_prompt: None,
config: RequestConfig::default(),
history: vec![],
}]);
store.append(id, &entries[0]).await.unwrap();
// Write a trace entry
let trace = TraceEntry {
ts: 1500,
turn: 0,
event: llm_worker::llm_client::event::Event::Ping(
llm_worker::llm_client::event::PingEvent { timestamp: None },
),
};
store.append_trace(id, &trace).await.unwrap();
// Log should have 1 entry, unaffected by trace
let log = store.read_all(id).await.unwrap();
assert_eq!(log.len(), 1);
// Trace file should exist separately
let trace_path = dir.path().join(format!("{id}.trace.jsonl"));
assert!(trace_path.exists());
}
#[tokio::test]
async fn read_head_hash_returns_last_entry_hash() {
let dir = tempfile::tempdir().unwrap();
let store = FsStore::new(dir.path()).await.unwrap();
let id = new_session_id();
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"),
},
]);
for entry in &entries {
store.append(id, entry).await.unwrap();
}
let head = store.read_head_hash(id).await.unwrap();
assert_eq!(head.as_ref(), Some(&entries[1].hash));
}
+539
View File
@@ -0,0 +1,539 @@
mod common;
use std::sync::Arc;
use async_trait::async_trait;
use common::MockLlmClient;
use llm_worker::interceptor::{Interceptor, TurnEndAction};
use llm_worker::llm_client::event::{Event, ResponseStatus, StatusEvent};
use llm_worker::llm_client::types::{Item, RequestConfig};
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use llm_worker::Worker;
use session_store::{
EntryHash, FsStore, LogEntry, Outcome, SessionStartState, Store, collect_state,
};
// =============================================================================
// Helpers
// =============================================================================
fn simple_text_events() -> Vec<Event> {
vec![
Event::text_block_start(0),
Event::text_delta(0, "Hello!"),
Event::text_block_stop(0, None),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
}
fn tool_call_events() -> Vec<Vec<Event>> {
vec![
// 1st response: tool call
vec![
Event::tool_use_start(0, "call_1", "get_weather"),
Event::tool_input_delta(0, r#"{"city":"Tokyo"}"#),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
// 2nd response: final text
vec![
Event::text_block_start(0),
Event::text_delta(0, "It's sunny in Tokyo!"),
Event::text_block_stop(0, None),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
]
}
#[derive(Clone)]
struct MockWeatherTool;
#[async_trait]
impl Tool for MockWeatherTool {
async fn execute(&self, _input_json: &str) -> Result<ToolOutput, ToolError> {
Ok("Sunny, 25C".to_string().into())
}
}
fn weather_tool_definition() -> ToolDefinition {
Arc::new(|| {
let meta = ToolMeta::new("get_weather")
.description("Get weather")
.input_schema(serde_json::json!({
"type": "object",
"properties": {
"city": { "type": "string" }
},
"required": ["city"]
}));
(meta, Arc::new(MockWeatherTool) as Arc<dyn Tool>)
})
}
/// Policy that forces Pause on every turn end.
struct PausePolicy;
#[async_trait]
impl Interceptor for PausePolicy {
async fn on_turn_end(&self, _history: &[Item]) -> TurnEndAction {
TurnEndAction::Pause
}
}
async fn make_store() -> (tempfile::TempDir, FsStore) {
let dir = tempfile::tempdir().unwrap();
let store = FsStore::new(dir.path()).await.unwrap();
(dir, store)
}
/// Run a worker turn and persist via session-store functions.
/// Takes ownership of the worker (needed for lock/unlock) and returns it.
async fn run_and_persist(
worker: Worker<MockLlmClient>,
store: &FsStore,
session_id: session_store::SessionId,
head_hash: &mut Option<EntryHash>,
input: &str,
) -> (Worker<MockLlmClient>, llm_worker::WorkerResult) {
let history_before = worker.history().len();
let mut locked = worker.lock();
let result = locked.run(input).await;
let worker = locked.unlock();
let new_items = &worker.history()[history_before..];
session_store::save_delta(store, session_id, head_hash, new_items)
.await
.unwrap();
session_store::save_turn_end(store, session_id, head_hash, worker.turn_count())
.await
.unwrap();
let outcome = match &result {
Ok(llm_worker::WorkerResult::Finished) => Outcome::Finished,
Ok(llm_worker::WorkerResult::Paused) => Outcome::Paused,
Ok(llm_worker::WorkerResult::LimitReached) => Outcome::LimitReached,
Err(e) => Outcome::Error {
message: e.to_string(),
},
};
session_store::save_outcome(store, session_id, head_hash, outcome, worker.last_run_interrupted())
.await
.unwrap();
let r = result.unwrap();
(worker, r)
}
// =============================================================================
// Tests
// =============================================================================
#[tokio::test]
async fn session_run_logs_entries() {
let (_dir, store) = make_store().await;
let client = MockLlmClient::new(simple_text_events());
let worker = Worker::new(client);
let (sid, head_hash) = session_store::create_session(
&store,
SessionStartState {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: worker.history(),
},
)
.await
.unwrap();
let mut head_hash = Some(head_hash);
let (worker, _) = run_and_persist(worker, &store, sid, &mut head_hash, "Hi").await;
let _ = &worker;
let entries = store.read_all(sid).await.unwrap();
// SessionStart, UserInput, AssistantItems, TurnEnd, RunOutcome (at minimum)
assert!(
entries.len() >= 4,
"expected at least 4 entries, got {}",
entries.len()
);
// First entry is SessionStart
assert!(matches!(&entries[0].entry, LogEntry::SessionStart { .. }));
// Has a RunOutcome with Finished
let has_finished = entries.iter().any(|e| {
matches!(
&e.entry,
LogEntry::RunOutcome {
outcome: Outcome::Finished,
..
}
)
});
assert!(has_finished, "should have a Finished outcome");
// Verify hash chain integrity
assert!(entries[0].prev_hash.is_none());
for i in 1..entries.len() {
assert_eq!(
entries[i].prev_hash.as_ref(),
Some(&entries[i - 1].hash),
"hash chain broken at entry {}",
i
);
}
}
#[tokio::test]
async fn session_restore_round_trip() {
let (_dir, store) = make_store().await;
let client = MockLlmClient::new(simple_text_events());
let mut worker = Worker::new(client);
worker.set_system_prompt("You are helpful.");
let (sid, head_hash) = session_store::create_session(
&store,
SessionStartState {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: worker.history(),
},
)
.await
.unwrap();
let mut head_hash = Some(head_hash);
let (worker, _) = run_and_persist(worker, &store, sid, &mut head_hash, "Hi").await;
let original_history_len = worker.history().len();
let original_turn_count = worker.turn_count();
// Restore
let state = session_store::restore(&store, sid).await.unwrap();
assert_eq!(state.history.len(), original_history_len);
assert_eq!(state.turn_count, original_turn_count);
assert_eq!(state.system_prompt.as_deref(), Some("You are helpful."));
assert_eq!(state.head_hash, head_hash);
}
#[tokio::test]
async fn session_run_with_tool_call() {
let (_dir, store) = make_store().await;
let client = MockLlmClient::with_responses(tool_call_events());
let mut worker = Worker::new(client);
worker.register_tool(weather_tool_definition());
let (sid, head_hash) = session_store::create_session(
&store,
SessionStartState {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: worker.history(),
},
)
.await
.unwrap();
let mut head_hash = Some(head_hash);
let (_worker, _) = run_and_persist(worker, &store, sid, &mut head_hash, "What's the weather?").await;
let entries = store.read_all(sid).await.unwrap();
let has_tool_results = entries
.iter()
.any(|e| matches!(&e.entry, LogEntry::ToolResults { .. }));
assert!(has_tool_results, "should have ToolResults entry");
let has_assistant = entries
.iter()
.any(|e| matches!(&e.entry, LogEntry::AssistantItems { .. }));
assert!(has_assistant, "should have AssistantItems entry");
}
#[tokio::test]
async fn session_resume_after_pause() {
let (_dir, store) = make_store().await;
// First run: tool call with pause policy → Paused
let client = MockLlmClient::with_responses(tool_call_events());
let mut worker = Worker::new(client);
worker.register_tool(weather_tool_definition());
worker.set_interceptor(PausePolicy);
let (sid, head_hash) = session_store::create_session(
&store,
SessionStartState {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: worker.history(),
},
)
.await
.unwrap();
let mut head_hash = Some(head_hash);
let (_worker, result) = run_and_persist(worker, &store, sid, &mut head_hash, "Weather?").await;
assert!(matches!(result, llm_worker::WorkerResult::Paused));
// Check RunOutcome is Paused
let entries = store.read_all(sid).await.unwrap();
let has_paused = entries.iter().any(|e| {
matches!(
&e.entry,
LogEntry::RunOutcome {
outcome: Outcome::Paused,
..
}
)
});
assert!(has_paused, "should have Paused outcome");
// Restore state and verify
let state = session_store::restore(&store, sid).await.unwrap();
assert!(state.last_run_interrupted);
}
#[tokio::test]
async fn session_fork_preserves_state() {
let (_dir, store) = make_store().await;
let client = MockLlmClient::new(simple_text_events());
let mut worker = Worker::new(client);
worker.set_system_prompt("System prompt");
let (sid, head_hash) = session_store::create_session(
&store,
SessionStartState {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: worker.history(),
},
)
.await
.unwrap();
let mut head_hash = Some(head_hash);
let (worker, _) = run_and_persist(worker, &store, sid, &mut head_hash, "Hello").await;
let original_history_len = worker.history().len();
let fork_id = session_store::fork(
&store,
SessionStartState {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: worker.history(),
},
)
.await
.unwrap();
// Fork should have a SessionStart with the current history
let fork_entries = store.read_all(fork_id).await.unwrap();
assert_eq!(fork_entries.len(), 1);
assert!(matches!(
&fork_entries[0].entry,
LogEntry::SessionStart { .. }
));
let fork_state = collect_state(&fork_entries);
assert_eq!(fork_state.history.len(), original_history_len);
assert_eq!(fork_state.system_prompt.as_deref(), Some("System prompt"));
}
#[tokio::test]
async fn session_fork_at_truncates() {
let (_dir, store) = make_store().await;
let client = MockLlmClient::new(simple_text_events());
let worker = Worker::new(client);
let (sid, head_hash) = session_store::create_session(
&store,
SessionStartState {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: worker.history(),
},
)
.await
.unwrap();
let mut head_hash = Some(head_hash);
let (_worker, _) = run_and_persist(worker, &store, sid, &mut head_hash, "Hello").await;
let all_entries = store.read_all(sid).await.unwrap();
assert!(all_entries.len() > 2);
// Fork at the hash of the 2nd entry (SessionStart + UserInput)
let at_hash = &all_entries[1].hash;
let fork_id = session_store::fork_at(&store, sid, at_hash).await.unwrap();
let fork_entries = store.read_all(fork_id).await.unwrap();
assert_eq!(fork_entries.len(), 1); // Just the new SessionStart
let fork_state = collect_state(&fork_entries);
// Should have the state from replaying only the first 2 entries
let original_truncated_state = collect_state(&all_entries[..2]);
assert_eq!(
fork_state.history.len(),
original_truncated_state.history.len()
);
}
#[tokio::test]
async fn session_config_changed_logged() {
let (_dir, store) = make_store().await;
let client = MockLlmClient::new(vec![]);
let mut worker = Worker::new(client);
let (sid, head_hash) = session_store::create_session(
&store,
SessionStartState {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: worker.history(),
},
)
.await
.unwrap();
let mut head_hash = Some(head_hash);
// Modify config and log it
let new_config = RequestConfig::default().with_temperature(0.7);
worker.set_request_config(new_config.clone());
session_store::save_config_changed(&store, sid, &mut head_hash, &new_config)
.await
.unwrap();
let entries = store.read_all(sid).await.unwrap();
let has_config_changed = entries.iter().any(|e| {
matches!(
&e.entry,
LogEntry::ConfigChanged { config, .. } if config.temperature == Some(0.7)
)
});
assert!(has_config_changed, "should have ConfigChanged entry");
}
#[tokio::test]
async fn session_cache_lock_unlock_logged() {
let (_dir, store) = make_store().await;
let client = MockLlmClient::new(vec![]);
let worker = Worker::new(client);
let (sid, head_hash) = session_store::create_session(
&store,
SessionStartState {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: worker.history(),
},
)
.await
.unwrap();
let mut head_hash = Some(head_hash);
session_store::save_cache_locked(&store, sid, &mut head_hash, 5)
.await
.unwrap();
session_store::save_cache_unlocked(&store, sid, &mut head_hash)
.await
.unwrap();
let entries = store.read_all(sid).await.unwrap();
let has_locked = entries.iter().any(|e| {
matches!(
&e.entry,
LogEntry::Locked {
locked_prefix_len: 5,
..
}
)
});
assert!(has_locked, "should have Locked entry");
let has_unlocked = entries
.iter()
.any(|e| matches!(&e.entry, LogEntry::CacheUnlocked { .. }));
assert!(has_unlocked, "should have CacheUnlocked entry");
// State after all entries: unlocked
let state = collect_state(&entries);
assert_eq!(state.locked_prefix_len, 0);
}
#[tokio::test]
async fn session_auto_forks_on_conflict() {
let (_dir, store) = make_store().await;
// Create a session
let client_a = MockLlmClient::new(simple_text_events());
let worker_a = Worker::new(client_a);
let (original_sid, head_hash) = session_store::create_session(
&store,
SessionStartState {
system_prompt: worker_a.get_system_prompt(),
config: worker_a.request_config(),
history: worker_a.history(),
},
)
.await
.unwrap();
let mut session_id = original_sid;
let mut head_hash = Some(head_hash);
// Simulate another Pod writing to the same session behind our back
let extra_entry = LogEntry::UserInput {
ts: 9999,
item: Item::user_message("Interloper"),
};
let current_head = store.read_head_hash(original_sid).await.unwrap();
let hash = session_store::compute_hash(current_head.as_ref(), &extra_entry);
let hashed = session_store::HashedEntry {
hash,
prev_hash: current_head,
entry: extra_entry,
};
store.append(original_sid, &hashed).await.unwrap();
// Now head_hash is stale — ensure_head_or_fork should auto-fork
session_store::ensure_head_or_fork(
&store,
&mut session_id,
&mut head_hash,
SessionStartState {
system_prompt: worker_a.get_system_prompt(),
config: worker_a.request_config(),
history: worker_a.history(),
},
)
.await
.unwrap();
// session_id should now be different
assert_ne!(session_id, original_sid);
// The fork session should exist and have entries
let fork_entries = store.read_all(session_id).await.unwrap();
assert!(!fork_entries.is_empty());
// Original session should still have the interloper entry
let original_entries = store.read_all(original_sid).await.unwrap();
let has_interloper = original_entries.iter().any(|e| {
if let LogEntry::UserInput { item, .. } = &e.entry {
item.is_user_message()
} else {
false
}
});
assert!(has_interloper);
}