use std::collections::{BTreeMap, HashMap};
use std::path::Path;
use std::sync::Mutex;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::{Digest, Sha256};
use crate::agent::AgentOutput;
use crate::error::Error;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ResumeMode {
Fresh,
Resume,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub(crate) struct CallKey {
pub content_hash: String,
pub occurrence: u64,
}
#[derive(Debug, Serialize, Deserialize)]
struct JournalRecord {
v: u16,
key: CallKey,
output: AgentOutput,
}
#[derive(Debug)]
pub struct RunJournal {
entries: Mutex<HashMap<CallKey, AgentOutput>>,
seen: Mutex<HashMap<String, u64>>,
file: Mutex<std::fs::File>,
}
impl RunJournal {
pub fn open(path: &Path, mode: ResumeMode) -> Result<Self, Error> {
use std::fs::OpenOptions;
use std::io::{BufRead, BufReader};
if let Some(parent) = path.parent()
&& !parent.as_os_str().is_empty()
{
std::fs::create_dir_all(parent).map_err(|e| Error::Store(e.to_string()))?;
}
let mut entries: HashMap<CallKey, AgentOutput> = HashMap::new();
if mode == ResumeMode::Resume && path.exists() {
let file = std::fs::File::open(path).map_err(|e| Error::Store(e.to_string()))?;
for line in BufReader::new(file).lines() {
let line = line.map_err(|e| Error::Store(e.to_string()))?;
if line.trim().is_empty() {
continue;
}
if let Ok(rec) = serde_json::from_str::<JournalRecord>(&line) {
entries.insert(rec.key, rec.output);
}
}
}
let file = match mode {
ResumeMode::Fresh => OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(path),
ResumeMode::Resume => OpenOptions::new().append(true).create(true).open(path),
}
.map_err(|e| Error::Store(e.to_string()))?;
Ok(Self {
entries: Mutex::new(entries),
seen: Mutex::new(HashMap::new()),
file: Mutex::new(file),
})
}
pub(crate) fn next_occurrence(&self, content_hash: &str) -> u64 {
let mut seen = self.seen.lock().expect("journal seen lock poisoned");
let counter = seen.entry(content_hash.to_string()).or_insert(0);
let occurrence = *counter;
*counter += 1;
occurrence
}
pub(crate) fn lookup(&self, key: &CallKey) -> Option<AgentOutput> {
self.entries
.lock()
.expect("journal entries lock poisoned")
.get(key)
.cloned()
}
pub(crate) fn append(&self, key: &CallKey, output: &AgentOutput) -> Result<(), Error> {
use std::io::Write;
let record = JournalRecord {
v: 1,
key: key.clone(),
output: output.clone(),
};
let mut line = serde_json::to_string(&record)?;
line.push('\n');
{
let mut file = self.file.lock().expect("journal file lock poisoned");
file.write_all(line.as_bytes())
.map_err(|e| Error::Store(e.to_string()))?;
file.flush().map_err(|e| Error::Store(e.to_string()))?;
}
self.entries
.lock()
.expect("journal entries lock poisoned")
.insert(key.clone(), output.clone());
Ok(())
}
}
pub(crate) fn canonical_json(value: &Value) -> Value {
match value {
Value::Object(map) => {
let sorted: BTreeMap<String, Value> = map
.iter()
.map(|(k, v)| (k.clone(), canonical_json(v)))
.collect();
Value::Object(sorted.into_iter().collect())
}
Value::Array(items) => Value::Array(items.iter().map(canonical_json).collect()),
other => other.clone(),
}
}
fn to_hex(bytes: &[u8]) -> String {
use std::fmt::Write;
let mut s = String::with_capacity(bytes.len() * 2);
for b in bytes {
let _ = write!(s, "{b:02x}");
}
s
}
pub(crate) fn content_hash(prompt: &str, model: Option<&str>, schema: Option<&Value>) -> String {
let mut hasher = Sha256::new();
hasher.update(prompt.as_bytes());
hasher.update([0u8]);
hasher.update(model.unwrap_or("").as_bytes());
hasher.update([0u8]);
if let Some(schema) = schema {
let canonical = serde_json::to_string(&canonical_json(schema)).unwrap_or_default();
hasher.update(canonical.as_bytes());
}
to_hex(&hasher.finalize())
}
pub fn derive_run_id(seed: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(seed.as_bytes());
to_hex(&hasher.finalize())
}
#[cfg(test)]
mod tests {
use super::*;
fn out(text: &str, input: u32, output: u32) -> AgentOutput {
AgentOutput {
result: text.to_string(),
tokens_used: crate::llm::types::TokenUsage {
input_tokens: input,
output_tokens: output,
..Default::default()
},
..Default::default()
}
}
fn key(hash: &str, occurrence: u64) -> CallKey {
CallKey {
content_hash: hash.to_string(),
occurrence,
}
}
#[test]
fn canonical_json_sorts_nested_keys() {
let a = serde_json::json!({ "b": 1, "a": { "y": 2, "x": 3 } });
let b = serde_json::json!({ "a": { "x": 3, "y": 2 }, "b": 1 });
assert_eq!(canonical_json(&a), canonical_json(&b));
let s = serde_json::to_string(&canonical_json(&a)).unwrap();
assert!(s.starts_with(r#"{"a":"#), "not sorted: {s}");
}
#[test]
fn content_hash_stable_under_schema_key_reorder() {
let s1 = serde_json::json!({ "type": "object", "required": ["x"] });
let s2 = serde_json::json!({ "required": ["x"], "type": "object" });
assert_eq!(
content_hash("p", None, Some(&s1)),
content_hash("p", None, Some(&s2)),
"reordered-but-equal schemas must hash the same"
);
}
#[test]
fn content_hash_distinguishes_inputs() {
let base = content_hash("prompt", None, None);
assert_ne!(base, content_hash("other", None, None), "prompt matters");
assert_ne!(
base,
content_hash("prompt", Some("haiku"), None),
"model matters"
);
let schema = serde_json::json!({ "type": "string" });
assert_ne!(
base,
content_hash("prompt", None, Some(&schema)),
"schema matters"
);
}
#[test]
fn derive_run_id_is_deterministic_and_seed_sensitive() {
assert_eq!(derive_run_id("a|b"), derive_run_id("a|b"));
assert_ne!(derive_run_id("a|b"), derive_run_id("a|c"));
}
#[test]
fn fresh_journal_is_empty() {
let dir = tempfile::tempdir().unwrap();
let j = RunJournal::open(&dir.path().join("j.jsonl"), ResumeMode::Fresh).unwrap();
assert!(j.lookup(&key("h", 0)).is_none());
}
#[test]
fn append_then_lookup_roundtrips() {
let dir = tempfile::tempdir().unwrap();
let j = RunJournal::open(&dir.path().join("j.jsonl"), ResumeMode::Fresh).unwrap();
j.append(&key("h", 0), &out("hello", 10, 5)).unwrap();
let got = j.lookup(&key("h", 0)).expect("hit");
assert_eq!(got.result, "hello");
assert_eq!(got.tokens_used.input_tokens, 10);
}
#[test]
fn occurrence_disambiguates_identical_content() {
let dir = tempfile::tempdir().unwrap();
let j = RunJournal::open(&dir.path().join("j.jsonl"), ResumeMode::Fresh).unwrap();
assert_eq!(j.next_occurrence("h"), 0);
assert_eq!(j.next_occurrence("h"), 1);
assert_eq!(j.next_occurrence("other"), 0);
j.append(&key("h", 0), &out("first", 1, 1)).unwrap();
j.append(&key("h", 1), &out("second", 1, 1)).unwrap();
assert_eq!(j.lookup(&key("h", 0)).unwrap().result, "first");
assert_eq!(j.lookup(&key("h", 1)).unwrap().result, "second");
}
#[test]
fn resume_reads_existing_entries() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("j.jsonl");
{
let j = RunJournal::open(&path, ResumeMode::Fresh).unwrap();
j.append(&key("h", 0), &out("persisted", 7, 3)).unwrap();
}
let j2 = RunJournal::open(&path, ResumeMode::Resume).unwrap();
assert_eq!(j2.lookup(&key("h", 0)).unwrap().result, "persisted");
}
#[test]
fn fresh_truncates_existing() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("j.jsonl");
{
let j = RunJournal::open(&path, ResumeMode::Fresh).unwrap();
j.append(&key("h", 0), &out("old", 1, 1)).unwrap();
}
let j2 = RunJournal::open(&path, ResumeMode::Fresh).unwrap();
assert!(j2.lookup(&key("h", 0)).is_none());
}
#[test]
fn torn_final_line_is_skipped_on_resume() {
use std::io::Write;
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("j.jsonl");
{
let j = RunJournal::open(&path, ResumeMode::Fresh).unwrap();
j.append(&key("h", 0), &out("good", 1, 1)).unwrap();
}
{
let mut f = std::fs::OpenOptions::new()
.append(true)
.open(&path)
.unwrap();
f.write_all(br#"{"v":1,"key":{"content_hash":"h","occ"#)
.unwrap();
}
let j2 = RunJournal::open(&path, ResumeMode::Resume).unwrap();
assert_eq!(j2.lookup(&key("h", 0)).unwrap().result, "good");
}
#[test]
fn resume_on_missing_file_is_empty_not_error() {
let dir = tempfile::tempdir().unwrap();
let j = RunJournal::open(&dir.path().join("absent.jsonl"), ResumeMode::Resume).unwrap();
assert!(j.lookup(&key("h", 0)).is_none());
}
}