Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix(session): preflight migration and retain timestamps

Browse Source
This commit is contained in:
Hare
2026-08-30 09:42:46 +09:00
parent 89e6a6215a
commit f8a7c46cf9
10 changed files with 223 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/session-store/src/public_snapshot.rs
+31 -14
View File
@@ -33,6 +33,7 @@ pub fn project_session_snapshot(session_id: SessionId, log: &[LogEntry]) -> Sess
for (log_index, record) in log.iter().enumerate() {
match record {
LogEntry::SegmentStart {
ts,
session_id,
history,
..
@@ -41,57 +42,65 @@ pub fn project_session_snapshot(session_id: SessionId, log: &[LogEntry]) -> Sess
entries.clear();
for (item_index, item) in history.iter().enumerate() {
if let Some(data) = project_item(item) {
entries.push(legacy_entry(&session_key, log_index, item_index, data));
entries.push(legacy_entry(&session_key, log_index, item_index, *ts, data));
}
}
}
LogEntry::AnnotatedSegmentStart {
ts,
session_id,
history,
..
} => {
session_key = *session_id;
entries.clear();
extend_history(&mut entries, history, None);
extend_history(&mut entries, history, None, *ts);
}
LogEntry::UserInput { segments, .. } => entries.push(legacy_entry(
LogEntry::UserInput { ts, segments, .. } => entries.push(legacy_entry(
&session_key,
log_index,
0,
*ts,
SessionSnapshotEntryData::UserInput {
segments: segments.clone(),
},
)),
LogEntry::AnnotatedUserInput {
segments, history, ..
} => extend_history(&mut entries, history, Some(segments)),
LogEntry::AssistantItem { item, .. } | LogEntry::ToolResult { item, .. } => {
ts,
segments,
history,
..
} => extend_history(&mut entries, history, Some(segments), *ts),
LogEntry::AssistantItem { ts, item } | LogEntry::ToolResult { ts, item } => {
if let Some(data) = project_item(item) {
entries.push(legacy_entry(&session_key, log_index, 0, data));
entries.push(legacy_entry(&session_key, log_index, 0, *ts, data));
}
}
LogEntry::AnnotatedAssistantItem { entry, .. }
| LogEntry::AnnotatedToolResult { entry, .. } => {
LogEntry::AnnotatedAssistantItem { ts, entry }
| LogEntry::AnnotatedToolResult { ts, entry } => {
if let Some(data) = project_item(&entry.item) {
entries.push(history_entry(entry, data));
entries.push(history_entry(entry, *ts, data));
}
}
LogEntry::SystemItem { item, .. } => entries.push(system_entry(
LogEntry::SystemItem { ts, item } => entries.push(system_entry(
item,
legacy_entry_id(&session_key, log_index, 0),
*ts,
SessionEntryProvenance::LegacyUnknown,
Vec::new(),
)),
LogEntry::AnnotatedSystemItem { entry, .. } => entries.push(system_entry(
LogEntry::AnnotatedSystemItem { ts, entry } => entries.push(system_entry(
&entry.item,
entry.metadata.entry_id.0.clone(),
*ts,
provenance(&entry.metadata.origin),
derivation_ids(entry),
)),
LogEntry::RunErrored { message, .. } => entries.push(legacy_entry(
LogEntry::RunErrored { ts, message, .. } => entries.push(legacy_entry(
&session_key,
log_index,
0,
*ts,
SessionSnapshotEntryData::RunError {
message: message.clone(),
},
@@ -116,6 +125,7 @@ fn extend_history(
output: &mut Vec<SessionSnapshotEntry>,
history: &[LoggedHistoryEntry],
input_segments: Option<&Vec<Segment>>,
timestamp: u64,
) {
let mut attached_segments = false;
for entry in history {
@@ -136,16 +146,18 @@ fn extend_history(
};
data
};
output.push(history_entry(entry, data));
output.push(history_entry(entry, timestamp, data));
}
}
fn history_entry(
entry: &LoggedHistoryEntry,
timestamp: u64,
data: SessionSnapshotEntryData,
) -> SessionSnapshotEntry {
SessionSnapshotEntry {
entry_id: entry.metadata.entry_id.0.clone(),
timestamp,
provenance: provenance(&entry.metadata.origin),
derived_from: entry
.metadata
@@ -182,10 +194,12 @@ fn legacy_entry(
session_key: &SessionId,
log_index: usize,
item_index: usize,
timestamp: u64,
data: SessionSnapshotEntryData,
) -> SessionSnapshotEntry {
SessionSnapshotEntry {
entry_id: legacy_entry_id(session_key, log_index, item_index),
timestamp,
provenance: SessionEntryProvenance::LegacyUnknown,
derived_from: Vec::new(),
data,
@@ -283,6 +297,7 @@ fn project_item(item: &LoggedItem) -> Option<SessionSnapshotEntryData> {
fn system_entry(
item: &SystemItem,
entry_id: String,
timestamp: u64,
provenance: SessionEntryProvenance,
derived_from: Vec<String>,
) -> SessionSnapshotEntry {
@@ -298,6 +313,7 @@ fn system_entry(
.to_owned();
SessionSnapshotEntry {
entry_id,
timestamp,
provenance,
derived_from,
data: SessionSnapshotEntryData::SystemItem {
@@ -351,6 +367,7 @@ mod tests {
let second = project_session_snapshot(session_id, &log);
assert_eq!(first, second);
assert_eq!(first.entries.len(), 1);
assert_eq!(first.entries[0].timestamp, 1);
assert_eq!(
first.entries[0].provenance,
SessionEntryProvenance::LegacyUnknown
crates/session-store/src/worker_session_store.rs
+86 -5
View File
@@ -381,6 +381,15 @@ fn segment_log_paths(root: &Path) -> Result<Vec<(SegmentId, PathBuf)>, StoreErro
}
fn migrate_segment_logs_to_v3(root: &Path, session_id: SessionId) -> Result<(), StoreError> {
struct MigrationPlan {
path: PathBuf,
source: Vec<u8>,
output: Vec<u8>,
}
// Phase 1 is strictly read-only. Every segment must parse and canonicalize
// successfully before the first authoritative byte is replaced.
let mut plans = Vec::new();
for (segment_id, path) in segment_log_paths(root)? {
let source = fs::read(&path)?;
let entries: Vec<LogEntry> = parse_jsonl(&source).map_err(|error| StoreError::Corrupt {
@@ -403,19 +412,30 @@ fn migrate_segment_logs_to_v3(root: &Path, session_id: SessionId) -> Result<(),
serde_json::to_writer(&mut output, &entry)?;
output.push(b'\n');
}
plans.push(MigrationPlan {
path,
source,
output,
});
}
// Opening a Session is the exclusive restore boundary, but retain an
// unchanged-source fence so a racing writer cannot be silently lost.
if fs::read(&path)? != source {
// Fence the complete preflight snapshot before starting phase 2. Session
// open is the exclusive restore boundary; this additionally fails closed
// if an unexpected writer raced the preflight.
for plan in &plans {
if fs::read(&plan.path)? != plan.source {
return Err(StoreError::Corrupt {
line: 0,
message: format!(
"Worker Session segment changed during migration: {}",
path.display()
plan.path.display()
),
});
}
atomic_write_bytes(&path, &output)?;
}
for plan in plans {
atomic_write_bytes(&plan.path, &plan.output)?;
}
Ok(())
}
@@ -774,12 +794,73 @@ mod tests {
&reopened.read_all(session_id, segment_id).unwrap(),
);
assert_eq!(snapshot.entries.len(), 3);
assert_eq!(
snapshot
.entries
.iter()
.map(|entry| entry.timestamp)
.collect::<Vec<_>>(),
vec![1, 2, 3]
);
assert!(snapshot.entries.iter().all(|entry| {
entry.provenance == protocol::SessionEntryProvenance::LegacyUnknown
&& entry.entry_id.len() <= 64
}));
}
#[test]
fn schema_v2_preflight_keeps_earlier_segments_unchanged_when_later_is_corrupt() {
let root = tempfile::tempdir().unwrap();
let session_id = new_session_id();
let valid_segment = uuid::Uuid::from_u128(1);
let corrupt_segment = uuid::Uuid::from_u128(2);
fs::create_dir_all(root.path().join(SEGMENTS_DIR)).unwrap();
atomic_write_json(
&root.path().join(SESSION_FILE),
&SessionManifest {
schema_version: PREVIOUS_SESSION_SCHEMA_VERSION,
session_id,
},
)
.unwrap();
let manifest_before = fs::read(root.path().join(SESSION_FILE)).unwrap();
let valid_path = root
.path()
.join(SEGMENTS_DIR)
.join(format!("{valid_segment}.jsonl"));
let valid_entry = LogEntry::SegmentStart {
ts: 1,
session_id,
system_prompt: None,
config: agen::llm_client::RequestConfig::default(),
history: vec![LoggedItem::from(agen::Item::assistant_message("prior"))],
forked_from: None,
compacted_from: None,
};
let mut valid_bytes = serde_json::to_vec(&valid_entry).unwrap();
valid_bytes.push(b'\n');
fs::write(&valid_path, &valid_bytes).unwrap();
let corrupt_path = root
.path()
.join(SEGMENTS_DIR)
.join(format!("{corrupt_segment}.jsonl"));
fs::write(&corrupt_path, b"{not-json}\n").unwrap();
let corrupt_before = fs::read(&corrupt_path).unwrap();
let error = match WorkerSessionStore::new(root.path()) {
Ok(_) => panic!("later corrupt segment must fail migration preflight"),
Err(error) => error,
};
assert!(matches!(error, StoreError::Corrupt { .. }));
assert_eq!(fs::read(&valid_path).unwrap(), valid_bytes);
assert_eq!(fs::read(&corrupt_path).unwrap(), corrupt_before);
assert_eq!(
fs::read(root.path().join(SESSION_FILE)).unwrap(),
manifest_before
);
}
#[test]
fn schema_v3_rejects_legacy_records_and_new_writes_are_canonical() {
let root = tempfile::tempdir().unwrap();