599 lines
20 KiB
Rust
599 lines
20 KiB
Rust
use agen::EngineResult;
|
|
use agen::llm_client::types::{Item, RequestConfig};
|
|
use session_store::{
|
|
FsStore, LogEntry, LoggedHistoryEntry, LoggedItem, LoggedSessionHistoryEntryId,
|
|
LoggedSessionHistoryMetadata, LoggedSessionHistoryOrigin, Store, TraceEntry, collect_state,
|
|
new_segment_id, new_session_id,
|
|
};
|
|
use std::io::Write;
|
|
|
|
fn annotated(item: Item) -> LoggedHistoryEntry {
|
|
LoggedHistoryEntry {
|
|
item: LoggedItem::from(item),
|
|
metadata: LoggedSessionHistoryMetadata {
|
|
input_metadata: None,
|
|
entry_id: LoggedSessionHistoryEntryId::new(),
|
|
origin: LoggedSessionHistoryOrigin::LegacyUnknown,
|
|
derivation: None,
|
|
},
|
|
}
|
|
}
|
|
|
|
fn nil_session_start(ts: u64, session_id: uuid::Uuid) -> LogEntry {
|
|
LogEntry::AnnotatedSegmentStart {
|
|
ts,
|
|
session_id,
|
|
system_prompt: None,
|
|
config: RequestConfig::default(),
|
|
history: vec![],
|
|
forked_from: None,
|
|
compacted_from: None,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn round_trip_write_and_read() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let store = FsStore::new(dir.path()).unwrap();
|
|
let sid = new_session_id();
|
|
let segid = new_segment_id();
|
|
|
|
let entries = vec![
|
|
LogEntry::AnnotatedSegmentStart {
|
|
ts: 1000,
|
|
session_id: sid,
|
|
system_prompt: Some("You are helpful.".into()),
|
|
config: RequestConfig::default().with_max_tokens(1024),
|
|
history: vec![],
|
|
forked_from: None,
|
|
compacted_from: None,
|
|
},
|
|
LogEntry::AnnotatedUserInput {
|
|
ts: 2000,
|
|
extensions: vec![],
|
|
segments: vec![protocol::Segment::text("Hello")],
|
|
history: vec![annotated(Item::user_message("Hello"))],
|
|
},
|
|
LogEntry::AnnotatedAssistantItem {
|
|
ts: 3000,
|
|
entry: annotated(Item::assistant_message("Hi there!")),
|
|
},
|
|
LogEntry::TurnEnd {
|
|
ts: 3100,
|
|
turn_count: 1,
|
|
},
|
|
LogEntry::RunCompleted {
|
|
entry_id: None,
|
|
ts: 3200,
|
|
interrupted: false,
|
|
result: EngineResult::Finished,
|
|
active_run_turn_count: None,
|
|
},
|
|
];
|
|
|
|
for entry in &entries {
|
|
store.append(sid, segid, entry).unwrap();
|
|
}
|
|
|
|
let read_back = store.read_all(sid, segid).unwrap();
|
|
assert_eq!(read_back.len(), entries.len());
|
|
|
|
let state = collect_state(&read_back);
|
|
assert_eq!(state.session_id, Some(sid));
|
|
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_eq!(state.entries_count, entries.len());
|
|
}
|
|
|
|
#[test]
|
|
fn create_segment_writes_all_entries() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let store = FsStore::new(dir.path()).unwrap();
|
|
let sid = new_session_id();
|
|
let segid = new_segment_id();
|
|
|
|
let entries = [LogEntry::AnnotatedSegmentStart {
|
|
ts: 1000,
|
|
session_id: sid,
|
|
system_prompt: None,
|
|
config: RequestConfig::default(),
|
|
history: vec![
|
|
annotated(Item::user_message("seed")),
|
|
annotated(Item::assistant_message("ok")),
|
|
],
|
|
forked_from: None,
|
|
compacted_from: None,
|
|
}];
|
|
|
|
store.create_segment(sid, segid, &entries).unwrap();
|
|
let read_back = store.read_all(sid, segid).unwrap();
|
|
assert_eq!(read_back.len(), 1);
|
|
|
|
let state = collect_state(&read_back);
|
|
assert_eq!(state.history.len(), 2);
|
|
assert_eq!(state.session_id, Some(sid));
|
|
}
|
|
|
|
#[test]
|
|
fn create_segment_replacement_keeps_open_old_log_and_commits_complete_new_log() {
|
|
use std::io::Read;
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let store = FsStore::new(dir.path()).unwrap();
|
|
let sid = new_session_id();
|
|
let segid = new_segment_id();
|
|
let old = vec![
|
|
nil_session_start(1, sid),
|
|
LogEntry::TurnEnd {
|
|
ts: 2,
|
|
turn_count: 1,
|
|
},
|
|
];
|
|
store.create_segment(sid, segid, &old).unwrap();
|
|
let path = dir
|
|
.path()
|
|
.join(sid.to_string())
|
|
.join(format!("{segid}.jsonl"));
|
|
let old_bytes = std::fs::read(&path).unwrap();
|
|
let mut old_reader = std::fs::File::open(&path).unwrap();
|
|
let new = vec![
|
|
nil_session_start(1, sid),
|
|
LogEntry::Extension {
|
|
ts: 3,
|
|
domain: "recovery-checkpoint".into(),
|
|
payload: serde_json::json!({"receipt": "retained"}),
|
|
},
|
|
];
|
|
store.create_segment(sid, segid, &new).unwrap();
|
|
let mut retained_inode = Vec::new();
|
|
old_reader.read_to_end(&mut retained_inode).unwrap();
|
|
// An in-place fs::write would truncate the inode held by this reader.
|
|
assert_eq!(retained_inode, old_bytes);
|
|
let reopened = FsStore::new(dir.path()).unwrap();
|
|
assert_eq!(
|
|
serde_json::to_value(reopened.read_all(sid, segid).unwrap()).unwrap(),
|
|
serde_json::to_value(&new).unwrap()
|
|
);
|
|
assert_eq!(reopened.read_entry_count(sid, segid).unwrap(), new.len());
|
|
assert_eq!(
|
|
std::fs::read_dir(path.parent().unwrap()).unwrap().count(),
|
|
2,
|
|
"successful replacement leaves only the segment and Session schema marker"
|
|
);
|
|
assert!(path.parent().unwrap().join("session-schema.json").is_file());
|
|
}
|
|
|
|
#[test]
|
|
fn create_segment_atomic_replacement_readers_only_observe_complete_old_or_new() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let store = FsStore::new(dir.path()).unwrap();
|
|
let sid = new_session_id();
|
|
let segid = new_segment_id();
|
|
let old = vec![
|
|
nil_session_start(1, sid),
|
|
LogEntry::Extension {
|
|
ts: 2,
|
|
domain: "old-evidence".into(),
|
|
payload: serde_json::json!({"result": "old"}),
|
|
},
|
|
];
|
|
let new = vec![
|
|
nil_session_start(1, sid),
|
|
LogEntry::Extension {
|
|
ts: 3,
|
|
domain: "new-checkpoint".into(),
|
|
payload: serde_json::json!({"result": "new"}),
|
|
},
|
|
];
|
|
store.create_segment(sid, segid, &old).unwrap();
|
|
let old_json = serde_json::to_value(&old).unwrap();
|
|
let new_json = serde_json::to_value(&new).unwrap();
|
|
let barrier = std::sync::Arc::new(std::sync::Barrier::new(2));
|
|
let reader_store = store.clone();
|
|
let reader_barrier = barrier.clone();
|
|
let reader = std::thread::spawn(move || {
|
|
reader_barrier.wait();
|
|
for _ in 0..1000 {
|
|
let observed =
|
|
serde_json::to_value(reader_store.read_all(sid, segid).unwrap()).unwrap();
|
|
assert!(
|
|
observed == old_json || observed == new_json,
|
|
"replacement must never expose a partial log: {observed}"
|
|
);
|
|
}
|
|
});
|
|
barrier.wait();
|
|
for _ in 0..50 {
|
|
store.create_segment(sid, segid, &new).unwrap();
|
|
store.create_segment(sid, segid, &old).unwrap();
|
|
}
|
|
reader.join().unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn list_sessions_and_segments() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let store = FsStore::new(dir.path()).unwrap();
|
|
|
|
let sid_a = new_session_id();
|
|
std::thread::sleep(std::time::Duration::from_millis(2));
|
|
let sid_b = new_session_id();
|
|
|
|
let seg_a1 = new_segment_id();
|
|
let seg_a2 = new_segment_id();
|
|
let seg_b1 = new_segment_id();
|
|
|
|
store
|
|
.append(sid_a, seg_a1, &nil_session_start(1, sid_a))
|
|
.unwrap();
|
|
store
|
|
.append(sid_a, seg_a2, &nil_session_start(2, sid_a))
|
|
.unwrap();
|
|
store
|
|
.append(sid_b, seg_b1, &nil_session_start(3, sid_b))
|
|
.unwrap();
|
|
|
|
let sessions = store.list_sessions().unwrap();
|
|
assert_eq!(sessions, vec![sid_b, sid_a]); // newest first
|
|
|
|
let segs_a = store.list_segments(sid_a).unwrap();
|
|
assert!(segs_a.contains(&seg_a1) && segs_a.contains(&seg_a2));
|
|
assert_eq!(segs_a.len(), 2);
|
|
|
|
let segs_b = store.list_segments(sid_b).unwrap();
|
|
assert_eq!(segs_b, vec![seg_b1]);
|
|
}
|
|
|
|
#[test]
|
|
fn exists_returns_correct_state() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let store = FsStore::new(dir.path()).unwrap();
|
|
let sid = new_session_id();
|
|
let segid = new_segment_id();
|
|
|
|
assert!(!store.exists(sid, segid).unwrap());
|
|
|
|
store
|
|
.append(sid, segid, &nil_session_start(1000, sid))
|
|
.unwrap();
|
|
|
|
assert!(store.exists(sid, segid).unwrap());
|
|
}
|
|
|
|
#[test]
|
|
fn not_found_error_for_missing_segment() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let store = FsStore::new(dir.path()).unwrap();
|
|
let sid = new_session_id();
|
|
let segid = new_segment_id();
|
|
|
|
let result = store.read_all(sid, segid);
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn trace_entries_in_separate_file() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let store = FsStore::new(dir.path()).unwrap();
|
|
let sid = new_session_id();
|
|
let segid = new_segment_id();
|
|
|
|
store
|
|
.append(sid, segid, &nil_session_start(1000, sid))
|
|
.unwrap();
|
|
|
|
let trace = TraceEntry {
|
|
ts: 1500,
|
|
turn: 0,
|
|
llm_call: Some(0),
|
|
payload: session_store::TracePayload::StreamEvent {
|
|
event: agen::llm_client::event::Event::Ping(agen::llm_client::event::PingEvent {
|
|
timestamp: None,
|
|
}),
|
|
},
|
|
};
|
|
store.append_trace(sid, segid, &trace).unwrap();
|
|
|
|
// Log should have 1 entry, unaffected by trace
|
|
let log = store.read_all(sid, segid).unwrap();
|
|
assert_eq!(log.len(), 1);
|
|
|
|
// Trace file should exist separately
|
|
let trace_path = dir
|
|
.path()
|
|
.join(sid.to_string())
|
|
.join(format!("{segid}.trace.jsonl"));
|
|
assert!(trace_path.exists());
|
|
}
|
|
|
|
#[test]
|
|
fn read_entry_count_matches_append_tally() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let store = FsStore::new(dir.path()).unwrap();
|
|
let sid = new_session_id();
|
|
let segid = new_segment_id();
|
|
|
|
let entries = [
|
|
LogEntry::AnnotatedSegmentStart {
|
|
ts: 1000,
|
|
session_id: sid,
|
|
system_prompt: None,
|
|
config: RequestConfig::default(),
|
|
history: vec![],
|
|
forked_from: None,
|
|
compacted_from: None,
|
|
},
|
|
LogEntry::AnnotatedUserInput {
|
|
ts: 2000,
|
|
extensions: vec![],
|
|
segments: vec![protocol::Segment::text("Hello")],
|
|
history: vec![annotated(Item::user_message("Hello"))],
|
|
},
|
|
];
|
|
|
|
for entry in &entries {
|
|
store.append(sid, segid, entry).unwrap();
|
|
}
|
|
|
|
assert_eq!(store.read_entry_count(sid, segid).unwrap(), entries.len());
|
|
}
|
|
|
|
#[test]
|
|
fn unterminated_utf8_tail_is_ignored_and_replaced_on_append() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let store = FsStore::new(dir.path()).unwrap();
|
|
let sid = new_session_id();
|
|
let segid = new_segment_id();
|
|
let path = dir
|
|
.path()
|
|
.join(sid.to_string())
|
|
.join(format!("{segid}.jsonl"));
|
|
|
|
store
|
|
.append(sid, segid, &nil_session_start(1, sid))
|
|
.unwrap();
|
|
std::fs::OpenOptions::new()
|
|
.append(true)
|
|
.open(&path)
|
|
.unwrap()
|
|
// First byte of a three-byte UTF-8 code point, matching an ENOSPC
|
|
// partial write observed in a real session log.
|
|
.write_all(&[0xe3])
|
|
.unwrap();
|
|
|
|
assert_eq!(store.read_all(sid, segid).unwrap().len(), 1);
|
|
assert_eq!(store.read_entry_count(sid, segid).unwrap(), 1);
|
|
|
|
let next = LogEntry::AnnotatedUserInput {
|
|
ts: 2,
|
|
extensions: vec![],
|
|
segments: vec![protocol::Segment::text("recovered")],
|
|
history: vec![annotated(Item::user_message("recovered"))],
|
|
};
|
|
store.append(sid, segid, &next).unwrap();
|
|
|
|
let bytes = std::fs::read(&path).unwrap();
|
|
assert!(std::str::from_utf8(&bytes).is_ok());
|
|
assert_eq!(store.read_all(sid, segid).unwrap().len(), 2);
|
|
assert_eq!(store.read_entry_count(sid, segid).unwrap(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn newline_terminated_invalid_utf8_is_reported_as_corruption() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let store = FsStore::new(dir.path()).unwrap();
|
|
let sid = new_session_id();
|
|
let segid = new_segment_id();
|
|
let path = dir
|
|
.path()
|
|
.join(sid.to_string())
|
|
.join(format!("{segid}.jsonl"));
|
|
|
|
store
|
|
.append(sid, segid, &nil_session_start(1, sid))
|
|
.unwrap();
|
|
std::fs::OpenOptions::new()
|
|
.append(true)
|
|
.open(path)
|
|
.unwrap()
|
|
.write_all(&[0xe3, b'\n'])
|
|
.unwrap();
|
|
|
|
assert!(matches!(
|
|
store.read_all(sid, segid),
|
|
Err(session_store::StoreError::Corrupt { line: 2, .. })
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn lookup_session_of_finds_owning_session() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let store = FsStore::new(dir.path()).unwrap();
|
|
let sid = new_session_id();
|
|
let segid = new_segment_id();
|
|
|
|
assert_eq!(store.lookup_session_of(segid).unwrap(), None);
|
|
|
|
store
|
|
.append(sid, segid, &nil_session_start(1, sid))
|
|
.unwrap();
|
|
|
|
assert_eq!(store.lookup_session_of(segid).unwrap(), Some(sid));
|
|
}
|
|
|
|
// Storage/projection component test: both production layouts, temporary originals,
|
|
// no Feature, model, runtime or Host authority is substituted into this boundary.
|
|
#[test]
|
|
fn generic_system_append_and_fork_preserve_metadata_without_public_payload_leak() {
|
|
use session_store::{
|
|
FeatureInputPayload, HistoryInputMetadata, HistoryInputUsePolicy, InputCompactionUse,
|
|
InputEvidenceUse, InputObservationUse, LoggedHistoryDerivation, LoggedSystemHistoryEntry,
|
|
PublicSystemItem, PublicSystemPresentation, SystemInputEnvelope, WorkerSessionStore,
|
|
};
|
|
for worker_layout in [false, true] {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let sid = new_session_id();
|
|
let segment = new_segment_id();
|
|
let fork = new_segment_id();
|
|
let store: Box<dyn Store> = if worker_layout {
|
|
Box::new(WorkerSessionStore::new(dir.path()).unwrap())
|
|
} else {
|
|
Box::new(FsStore::new(dir.path()).unwrap())
|
|
};
|
|
store
|
|
.create_segment(sid, segment, &[nil_session_start(1, sid)])
|
|
.unwrap();
|
|
let entry = LoggedSystemHistoryEntry::from_input(
|
|
LoggedSessionHistoryOrigin::Provided {
|
|
producer: "new-feature".into(),
|
|
},
|
|
Some(LoggedHistoryDerivation {
|
|
sources: vec![LoggedSessionHistoryEntryId("source-entry".into())],
|
|
}),
|
|
SystemInputEnvelope::new(
|
|
"private model body",
|
|
HistoryInputMetadata::new(
|
|
HistoryInputUsePolicy {
|
|
recording_material: InputEvidenceUse::Exclude,
|
|
memory_evidence: InputEvidenceUse::Exclude,
|
|
observation: InputObservationUse::Hidden,
|
|
compaction: InputCompactionUse::ReferenceInput,
|
|
},
|
|
Some(
|
|
FeatureInputPayload::encode(
|
|
"new-feature:context",
|
|
101,
|
|
&serde_json::json!({"private": "secret-detail"}),
|
|
)
|
|
.unwrap(),
|
|
),
|
|
Some(PublicSystemItem {
|
|
kind: "context".into(),
|
|
body: "explicit public body".into(),
|
|
presentation: PublicSystemPresentation::Context,
|
|
}),
|
|
)
|
|
.unwrap(),
|
|
)
|
|
.unwrap(),
|
|
);
|
|
store
|
|
.append(
|
|
sid,
|
|
segment,
|
|
&LogEntry::AnnotatedSystemItem {
|
|
ts: 2,
|
|
entry: entry.clone(),
|
|
extensions: vec![],
|
|
},
|
|
)
|
|
.unwrap();
|
|
let saved = store.read_all(sid, segment).unwrap();
|
|
let LogEntry::AnnotatedSystemItem {
|
|
entry: restored, ..
|
|
} = &saved[1]
|
|
else {
|
|
panic!("missing system input")
|
|
};
|
|
assert_eq!(restored.to_history_entry(), entry.to_history_entry());
|
|
assert_eq!(
|
|
collect_state(&saved).history[0].as_text(),
|
|
Some("private model body")
|
|
);
|
|
let mut fork_seed = nil_session_start(3, sid);
|
|
if let LogEntry::AnnotatedSegmentStart {
|
|
history,
|
|
forked_from,
|
|
..
|
|
} = &mut fork_seed
|
|
{
|
|
history.push(restored.to_history_entry());
|
|
*forked_from = Some(session_store::SegmentOrigin {
|
|
segment_id: segment,
|
|
at_turn_index: 0,
|
|
});
|
|
}
|
|
store.create_segment(sid, fork, &[fork_seed]).unwrap();
|
|
let fork_saved = store.read_all(sid, fork).unwrap();
|
|
let LogEntry::AnnotatedSegmentStart { history, .. } = &fork_saved[0] else {
|
|
unreachable!()
|
|
};
|
|
assert_eq!(history, &[entry.to_history_entry()]);
|
|
let original = session_store::public_snapshot::project_session_snapshot(sid, &saved);
|
|
let inherited = session_store::public_snapshot::project_session_snapshot(sid, &fork_saved);
|
|
let public = serde_json::to_string(&original.entries).unwrap();
|
|
assert!(public.contains("explicit public body"));
|
|
for private in ["private model body", "secret-detail", "new-feature:context"] {
|
|
assert!(!public.contains(private), "private field leaked: {private}");
|
|
}
|
|
assert_eq!(original.entries.len(), 1);
|
|
assert_eq!(inherited.entries.len(), 1);
|
|
assert_eq!(original.entries[0].entry_id, inherited.entries[0].entry_id);
|
|
assert_eq!(
|
|
original.entries[0].derived_from,
|
|
inherited.entries[0].derived_from
|
|
);
|
|
assert_eq!(
|
|
original.entries[0].provenance,
|
|
inherited.entries[0].provenance
|
|
);
|
|
assert_eq!(
|
|
serde_json::to_value(&original.entries[0].data).unwrap(),
|
|
serde_json::to_value(&inherited.entries[0].data).unwrap()
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn invalid_generic_system_input_rejects_append_create_and_external_json_without_changing_originals()
|
|
{
|
|
use session_store::{StoreError, SystemItem, WorkerSessionStore};
|
|
for worker_layout in [false, true] {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let sid = new_session_id();
|
|
let segment = new_segment_id();
|
|
let store: Box<dyn Store> = if worker_layout {
|
|
Box::new(WorkerSessionStore::new(dir.path()).unwrap())
|
|
} else {
|
|
Box::new(FsStore::new(dir.path()).unwrap())
|
|
};
|
|
let seed = nil_session_start(1, sid);
|
|
store
|
|
.create_segment(sid, segment, std::slice::from_ref(&seed))
|
|
.unwrap();
|
|
let path = if worker_layout {
|
|
dir.path().join("segments").join(format!("{segment}.jsonl"))
|
|
} else {
|
|
dir.path()
|
|
.join(sid.to_string())
|
|
.join(format!("{segment}.jsonl"))
|
|
};
|
|
let before = std::fs::read(&path).unwrap();
|
|
let invalid = LogEntry::AnnotatedSystemItem {
|
|
ts: 2,
|
|
extensions: vec![],
|
|
entry: session_store::LoggedSystemHistoryEntry {
|
|
item: SystemItem::Input {
|
|
body: "no policy".into(),
|
|
},
|
|
metadata: LoggedSessionHistoryMetadata::legacy_unknown(),
|
|
},
|
|
};
|
|
assert!(matches!(
|
|
store.append(sid, segment, &invalid),
|
|
Err(StoreError::InvalidInputMetadata(_))
|
|
));
|
|
assert!(matches!(
|
|
store.create_segment(sid, segment, &[seed, invalid.clone()]),
|
|
Err(StoreError::InvalidInputMetadata(_))
|
|
));
|
|
assert!(
|
|
serde_json::from_value::<LogEntry>(serde_json::to_value(invalid).unwrap()).is_err()
|
|
);
|
|
assert_eq!(std::fs::read(path).unwrap(), before);
|
|
assert_eq!(store.read_all(sid, segment).unwrap().len(), 1);
|
|
}
|
|
}
|