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/protocol/src/lib.rs
+4
View File
@@ -394,6 +394,8 @@ pub struct SessionSnapshotEntry {
/// Stable identity from durable history metadata, or a deterministic /// Stable identity from durable history metadata, or a deterministic
/// identity derived from the legacy segment and log position. /// identity derived from the legacy segment and log position.
pub entry_id: String, pub entry_id: String,
/// Timestamp copied from the durable log record that commits this entry.
pub timestamp: u64,
pub provenance: SessionEntryProvenance, pub provenance: SessionEntryProvenance,
#[serde(default, skip_serializing_if = "Vec::is_empty")] #[serde(default, skip_serializing_if = "Vec::is_empty")]
pub derived_from: Vec<String>, pub derived_from: Vec<String>,
@@ -1536,6 +1538,7 @@ mod tests {
session: SessionSnapshot { session: SessionSnapshot {
entries: vec![SessionSnapshotEntry { entries: vec![SessionSnapshotEntry {
entry_id: "entry-1".into(), entry_id: "entry-1".into(),
timestamp: 1,
provenance: SessionEntryProvenance::HumanInput, provenance: SessionEntryProvenance::HumanInput,
derived_from: Vec::new(), derived_from: Vec::new(),
data: SessionSnapshotEntryData::UserInput { data: SessionSnapshotEntryData::UserInput {
@@ -1565,6 +1568,7 @@ mod tests {
parsed["data"]["session"]["entries"][0]["kind"], parsed["data"]["session"]["entries"][0]["kind"],
"user_input" "user_input"
); );
assert_eq!(parsed["data"]["session"]["entries"][0]["timestamp"], 1);
assert_eq!(parsed["data"]["greeting"]["worker_name"], "test"); assert_eq!(parsed["data"]["greeting"]["worker_name"], "test");
assert_eq!(parsed["data"]["greeting"]["tools"][0], "Read"); assert_eq!(parsed["data"]["greeting"]["tools"][0], "Read");
assert_eq!(parsed["data"]["greeting"]["context_window"], 200_000); assert_eq!(parsed["data"]["greeting"]["context_window"], 200_000);
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() { for (log_index, record) in log.iter().enumerate() {
match record { match record {
LogEntry::SegmentStart { LogEntry::SegmentStart {
ts,
session_id, session_id,
history, history,
.. ..
@@ -41,57 +42,65 @@ pub fn project_session_snapshot(session_id: SessionId, log: &[LogEntry]) -> Sess
entries.clear(); entries.clear();
for (item_index, item) in history.iter().enumerate() { for (item_index, item) in history.iter().enumerate() {
if let Some(data) = project_item(item) { 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 { LogEntry::AnnotatedSegmentStart {
ts,
session_id, session_id,
history, history,
.. ..
} => { } => {
session_key = *session_id; session_key = *session_id;
entries.clear(); 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, &session_key,
log_index, log_index,
0, 0,
*ts,
SessionSnapshotEntryData::UserInput { SessionSnapshotEntryData::UserInput {
segments: segments.clone(), segments: segments.clone(),
}, },
)), )),
LogEntry::AnnotatedUserInput { LogEntry::AnnotatedUserInput {
segments, history, .. ts,
} => extend_history(&mut entries, history, Some(segments)), segments,
LogEntry::AssistantItem { item, .. } | LogEntry::ToolResult { item, .. } => { history,
..
} => extend_history(&mut entries, history, Some(segments), *ts),
LogEntry::AssistantItem { ts, item } | LogEntry::ToolResult { ts, item } => {
if let Some(data) = project_item(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::AnnotatedAssistantItem { ts, entry }
| LogEntry::AnnotatedToolResult { entry, .. } => { | LogEntry::AnnotatedToolResult { ts, entry } => {
if let Some(data) = project_item(&entry.item) { 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, item,
legacy_entry_id(&session_key, log_index, 0), legacy_entry_id(&session_key, log_index, 0),
*ts,
SessionEntryProvenance::LegacyUnknown, SessionEntryProvenance::LegacyUnknown,
Vec::new(), Vec::new(),
)), )),
LogEntry::AnnotatedSystemItem { entry, .. } => entries.push(system_entry( LogEntry::AnnotatedSystemItem { ts, entry } => entries.push(system_entry(
&entry.item, &entry.item,
entry.metadata.entry_id.0.clone(), entry.metadata.entry_id.0.clone(),
*ts,
provenance(&entry.metadata.origin), provenance(&entry.metadata.origin),
derivation_ids(entry), derivation_ids(entry),
)), )),
LogEntry::RunErrored { message, .. } => entries.push(legacy_entry( LogEntry::RunErrored { ts, message, .. } => entries.push(legacy_entry(
&session_key, &session_key,
log_index, log_index,
0, 0,
*ts,
SessionSnapshotEntryData::RunError { SessionSnapshotEntryData::RunError {
message: message.clone(), message: message.clone(),
}, },
@@ -116,6 +125,7 @@ fn extend_history(
output: &mut Vec<SessionSnapshotEntry>, output: &mut Vec<SessionSnapshotEntry>,
history: &[LoggedHistoryEntry], history: &[LoggedHistoryEntry],
input_segments: Option<&Vec<Segment>>, input_segments: Option<&Vec<Segment>>,
timestamp: u64,
) { ) {
let mut attached_segments = false; let mut attached_segments = false;
for entry in history { for entry in history {
@@ -136,16 +146,18 @@ fn extend_history(
}; };
data data
}; };
output.push(history_entry(entry, data)); output.push(history_entry(entry, timestamp, data));
} }
} }
fn history_entry( fn history_entry(
entry: &LoggedHistoryEntry, entry: &LoggedHistoryEntry,
timestamp: u64,
data: SessionSnapshotEntryData, data: SessionSnapshotEntryData,
) -> SessionSnapshotEntry { ) -> SessionSnapshotEntry {
SessionSnapshotEntry { SessionSnapshotEntry {
entry_id: entry.metadata.entry_id.0.clone(), entry_id: entry.metadata.entry_id.0.clone(),
timestamp,
provenance: provenance(&entry.metadata.origin), provenance: provenance(&entry.metadata.origin),
derived_from: entry derived_from: entry
.metadata .metadata
@@ -182,10 +194,12 @@ fn legacy_entry(
session_key: &SessionId, session_key: &SessionId,
log_index: usize, log_index: usize,
item_index: usize, item_index: usize,
timestamp: u64,
data: SessionSnapshotEntryData, data: SessionSnapshotEntryData,
) -> SessionSnapshotEntry { ) -> SessionSnapshotEntry {
SessionSnapshotEntry { SessionSnapshotEntry {
entry_id: legacy_entry_id(session_key, log_index, item_index), entry_id: legacy_entry_id(session_key, log_index, item_index),
timestamp,
provenance: SessionEntryProvenance::LegacyUnknown, provenance: SessionEntryProvenance::LegacyUnknown,
derived_from: Vec::new(), derived_from: Vec::new(),
data, data,
@@ -283,6 +297,7 @@ fn project_item(item: &LoggedItem) -> Option<SessionSnapshotEntryData> {
fn system_entry( fn system_entry(
item: &SystemItem, item: &SystemItem,
entry_id: String, entry_id: String,
timestamp: u64,
provenance: SessionEntryProvenance, provenance: SessionEntryProvenance,
derived_from: Vec<String>, derived_from: Vec<String>,
) -> SessionSnapshotEntry { ) -> SessionSnapshotEntry {
@@ -298,6 +313,7 @@ fn system_entry(
.to_owned(); .to_owned();
SessionSnapshotEntry { SessionSnapshotEntry {
entry_id, entry_id,
timestamp,
provenance, provenance,
derived_from, derived_from,
data: SessionSnapshotEntryData::SystemItem { data: SessionSnapshotEntryData::SystemItem {
@@ -351,6 +367,7 @@ mod tests {
let second = project_session_snapshot(session_id, &log); let second = project_session_snapshot(session_id, &log);
assert_eq!(first, second); assert_eq!(first, second);
assert_eq!(first.entries.len(), 1); assert_eq!(first.entries.len(), 1);
assert_eq!(first.entries[0].timestamp, 1);
assert_eq!( assert_eq!(
first.entries[0].provenance, first.entries[0].provenance,
SessionEntryProvenance::LegacyUnknown 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> { 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)? { for (segment_id, path) in segment_log_paths(root)? {
let source = fs::read(&path)?; let source = fs::read(&path)?;
let entries: Vec<LogEntry> = parse_jsonl(&source).map_err(|error| StoreError::Corrupt { 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)?; serde_json::to_writer(&mut output, &entry)?;
output.push(b'\n'); output.push(b'\n');
} }
plans.push(MigrationPlan {
path,
source,
output,
});
}
// Opening a Session is the exclusive restore boundary, but retain an // Fence the complete preflight snapshot before starting phase 2. Session
// unchanged-source fence so a racing writer cannot be silently lost. // open is the exclusive restore boundary; this additionally fails closed
if fs::read(&path)? != source { // if an unexpected writer raced the preflight.
for plan in &plans {
if fs::read(&plan.path)? != plan.source {
return Err(StoreError::Corrupt { return Err(StoreError::Corrupt {
line: 0, line: 0,
message: format!( message: format!(
"Worker Session segment changed during migration: {}", "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(()) Ok(())
} }
@@ -774,12 +794,73 @@ mod tests {
&reopened.read_all(session_id, segment_id).unwrap(), &reopened.read_all(session_id, segment_id).unwrap(),
); );
assert_eq!(snapshot.entries.len(), 3); 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| { assert!(snapshot.entries.iter().all(|entry| {
entry.provenance == protocol::SessionEntryProvenance::LegacyUnknown entry.provenance == protocol::SessionEntryProvenance::LegacyUnknown
&& entry.entry_id.len() <= 64 && 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] #[test]
fn schema_v3_rejects_legacy_records_and_new_writes_are_canonical() { fn schema_v3_rejects_legacy_records_and_new_writes_are_canonical() {
let root = tempfile::tempdir().unwrap(); let root = tempfile::tempdir().unwrap();
crates/worker-runtime/src/runtime.rs
+1
View File
@@ -4588,6 +4588,7 @@ mod tests {
session: protocol::SessionSnapshot { session: protocol::SessionSnapshot {
entries: vec![protocol::SessionSnapshotEntry { entries: vec![protocol::SessionSnapshotEntry {
entry_id: "restored-log-entry".to_owned(), entry_id: "restored-log-entry".to_owned(),
timestamp: 1,
provenance: protocol::SessionEntryProvenance::LegacyUnknown, provenance: protocol::SessionEntryProvenance::LegacyUnknown,
derived_from: Vec::new(), derived_from: Vec::new(),
data: protocol::SessionSnapshotEntryData::RunError { data: protocol::SessionSnapshotEntryData::RunError {
crates/worker/src/feature/builtin/memory_extract.rs
+23 -54
View File
@@ -185,12 +185,9 @@ impl Tool for StageMemoryCandidateTool {
})?); })?);
} }
if matches!(params.kind, CandidateKind::Preference) if matches!(params.kind, CandidateKind::Preference)
&& entries.iter().any(|entry| { && entries
!matches!( .iter()
entry.origin, .any(|entry| !matches!(entry.origin, protocol::SessionEntryProvenance::HumanInput))
crate::WorkerHistoryProvenance::HumanInput { .. }
)
})
{ {
return Err(ToolError::InvalidArgument( return Err(ToolError::InvalidArgument(
"preference candidates require exclusively HumanInput evidence; model, Worker, Flow, backend, derived, and legacy-unknown origins are not preference authority" "preference candidates require exclusively HumanInput evidence; model, Worker, Flow, backend, derived, and legacy-unknown origins are not preference authority"
@@ -324,10 +321,23 @@ fn evidence_kind(entry: &SessionEntryEvidence) -> EvidenceKind {
} }
} }
fn evidence_origin(origin: &crate::WorkerHistoryProvenance) -> EvidenceOrigin { fn evidence_origin(origin: &protocol::SessionEntryProvenance) -> EvidenceOrigin {
use crate::WorkerHistoryProvenance as Origin; use protocol::SessionEntryProvenance as Origin;
let mut evidence = EvidenceOrigin { let kind = match origin {
kind: EvidenceOriginKind::LegacyUnknown, Origin::HumanInput => EvidenceOriginKind::HumanInput,
Origin::WorkerInput => EvidenceOriginKind::WorkerInput,
Origin::FlowInstruction => EvidenceOriginKind::FlowInstruction,
Origin::BackendInstruction => EvidenceOriginKind::BackendInstruction,
Origin::ModelOutput => EvidenceOriginKind::ModelOutput,
Origin::ToolOutput => EvidenceOriginKind::ToolOutput,
Origin::DerivedSummary => EvidenceOriginKind::DerivedSummary,
Origin::LegacyUnknown => EvidenceOriginKind::LegacyUnknown,
};
EvidenceOrigin {
kind,
// The public SessionSnapshot intentionally excludes account, Worker,
// Runtime, and Flow internals. Preserve the authenticated origin class
// without inventing missing control-plane identity fields.
account_id: None, account_id: None,
workspace_id: None, workspace_id: None,
runtime_id: None, runtime_id: None,
@@ -335,46 +345,7 @@ fn evidence_origin(origin: &crate::WorkerHistoryProvenance) -> EvidenceOrigin {
flow_selector: None, flow_selector: None,
flow_definition_id: None, flow_definition_id: None,
flow_definition_revision: None, flow_definition_revision: None,
};
match origin {
Origin::HumanInput { account_id } => {
evidence.kind = EvidenceOriginKind::HumanInput;
evidence.account_id = Some(account_id.clone());
}
Origin::WorkerInput { actor } => {
evidence.kind = EvidenceOriginKind::WorkerInput;
evidence.workspace_id = actor.workspace_id.clone();
evidence.runtime_id = actor.runtime_id.clone();
evidence.worker_id = Some(actor.worker_id.clone());
}
Origin::FlowInstruction {
selector,
definition_id,
definition_revision,
..
} => {
evidence.kind = EvidenceOriginKind::FlowInstruction;
evidence.flow_selector = Some(selector.clone());
evidence.flow_definition_id = Some(definition_id.clone());
evidence.flow_definition_revision = Some(*definition_revision);
}
Origin::BackendInstruction { .. } => evidence.kind = EvidenceOriginKind::BackendInstruction,
Origin::ModelOutput { worker } => {
evidence.kind = EvidenceOriginKind::ModelOutput;
evidence.workspace_id = worker.workspace_id.clone();
evidence.runtime_id = worker.runtime_id.clone();
evidence.worker_id = Some(worker.worker_id.clone());
}
Origin::ToolOutput { worker } => {
evidence.kind = EvidenceOriginKind::ToolOutput;
evidence.workspace_id = worker.workspace_id.clone();
evidence.runtime_id = worker.runtime_id.clone();
evidence.worker_id = Some(worker.worker_id.clone());
}
Origin::DerivedSummary => evidence.kind = EvidenceOriginKind::DerivedSummary,
Origin::LegacyUnknown => evidence.kind = EvidenceOriginKind::LegacyUnknown,
} }
evidence
} }
fn staging_evidence(entry: &SessionEntryEvidence) -> StagingEvidence { fn staging_evidence(entry: &SessionEntryEvidence) -> StagingEvidence {
@@ -502,12 +473,10 @@ mod tests {
} }
#[test] #[test]
fn human_origin_projects_account_authority_into_evidence() { fn public_human_origin_preserves_class_without_inventing_account_authority() {
let origin = evidence_origin(&crate::WorkerHistoryProvenance::HumanInput { let origin = evidence_origin(&protocol::SessionEntryProvenance::HumanInput);
account_id: "account-1".into(),
});
assert_eq!(origin.kind, EvidenceOriginKind::HumanInput); assert_eq!(origin.kind, EvidenceOriginKind::HumanInput);
assert_eq!(origin.account_id.as_deref(), Some("account-1")); assert_eq!(origin.account_id, None);
} }
#[test] #[test]
crates/worker/src/feature/builtin/worker_observation.rs
+1
View File
@@ -794,6 +794,7 @@ mod tests {
}; };
Some(protocol::SessionSnapshotEntry { Some(protocol::SessionSnapshotEntry {
entry_id: format!("fake-{index:08}"), entry_id: format!("fake-{index:08}"),
timestamp: index as u64,
provenance: protocol::SessionEntryProvenance::LegacyUnknown, provenance: protocol::SessionEntryProvenance::LegacyUnknown,
derived_from: Vec::new(), derived_from: Vec::new(),
data, data,
crates/worker/src/session_capture.rs
+70 -73
View File
@@ -109,7 +109,7 @@ impl ToolPart {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct OverviewItem { pub(crate) struct OverviewItem {
pub id: SessionEntryRef, pub id: SessionEntryRef,
pub origin: WorkerHistoryProvenance, pub origin: SessionEntryProvenance,
pub entry_range: [u64; 2], pub entry_range: [u64; 2],
pub kind: ReferenceKind, pub kind: ReferenceKind,
pub label: String, pub label: String,
@@ -120,7 +120,7 @@ pub(crate) struct OverviewItem {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct ReferenceEntry { pub(crate) struct ReferenceEntry {
pub id: SessionEntryRef, pub id: SessionEntryRef,
pub origin: WorkerHistoryProvenance, pub origin: SessionEntryProvenance,
pub entry_range: [u64; 2], pub entry_range: [u64; 2],
pub kind: ReferenceKind, pub kind: ReferenceKind,
pub tool_part: Option<ToolPart>, pub tool_part: Option<ToolPart>,
@@ -146,7 +146,7 @@ pub(crate) struct SearchOptions {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct SearchHit { pub(crate) struct SearchHit {
pub id: SessionEntryRef, pub id: SessionEntryRef,
pub origin: WorkerHistoryProvenance, pub origin: SessionEntryProvenance,
pub kind: ReferenceKind, pub kind: ReferenceKind,
pub tool_part: Option<ToolPart>, pub tool_part: Option<ToolPart>,
pub tool_name: Option<String>, pub tool_name: Option<String>,
@@ -192,7 +192,7 @@ impl Default for ReadOptions {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct ReadEntry { pub(crate) struct ReadEntry {
pub id: SessionEntryRef, pub id: SessionEntryRef,
pub origin: WorkerHistoryProvenance, pub origin: SessionEntryProvenance,
pub kind: ReferenceKind, pub kind: ReferenceKind,
pub tool_part: Option<ToolPart>, pub tool_part: Option<ToolPart>,
pub tool_name: Option<String>, pub tool_name: Option<String>,
@@ -211,7 +211,7 @@ pub(crate) struct ReadResult {
pub(crate) struct SessionEntryEvidence { pub(crate) struct SessionEntryEvidence {
pub segment_id: String, pub segment_id: String,
pub entry_ref: SessionEntryRef, pub entry_ref: SessionEntryRef,
pub origin: WorkerHistoryProvenance, pub origin: SessionEntryProvenance,
pub entry_range: [u64; 2], pub entry_range: [u64; 2],
pub kind: ReferenceKind, pub kind: ReferenceKind,
pub tool_part: Option<ToolPart>, pub tool_part: Option<ToolPart>,
@@ -220,10 +220,17 @@ pub(crate) struct SessionEntryEvidence {
pub excerpt: String, pub excerpt: String,
} }
#[derive(Debug, Clone)]
struct CapturedHistoryEntry {
item: Item,
entry_id: session_store::LoggedSessionHistoryEntryId,
origin: SessionEntryProvenance,
}
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct SessionCapture { pub(crate) struct SessionCapture {
segment_id: String, segment_id: String,
entries: Arc<Vec<HistoryEntry<SessionHistoryMetadata>>>, entries: Arc<Vec<CapturedHistoryEntry>>,
overview: Vec<OverviewItem>, overview: Vec<OverviewItem>,
index: Vec<ReferenceEntry>, index: Vec<ReferenceEntry>,
} }
@@ -280,35 +287,49 @@ impl SessionCapture {
SessionSnapshotEntryData::SystemItem { .. } SessionSnapshotEntryData::SystemItem { .. }
| SessionSnapshotEntryData::RunError { .. } => return None, | SessionSnapshotEntryData::RunError { .. } => return None,
}; };
Some(HistoryEntry::new( Some(CapturedHistoryEntry {
item, item,
public_snapshot_metadata(entry.entry_id, entry.provenance), entry_id: session_store::LoggedSessionHistoryEntryId(entry.entry_id),
)) origin: entry.provenance,
})
}) })
.collect(); .collect();
Self::from_history_entries(segment_id, entries) Self::from_captured_entries(segment_id, entries)
} }
pub(crate) fn new(segment_id: impl Into<String>, items: Vec<Item>) -> Self { pub(crate) fn new(segment_id: impl Into<String>, items: Vec<Item>) -> Self {
let entries = items let entries = items
.into_iter() .into_iter()
.enumerate() .enumerate()
.map(|(index, item)| { .map(|(index, item)| CapturedHistoryEntry {
let mut metadata = SessionHistoryMetadata::legacy_unknown(); item,
metadata.entry_id = entry_id: session_store::LoggedSessionHistoryEntryId(format!("{index:08}")),
session_store::LoggedSessionHistoryEntryId(format!("{index:08}")); origin: SessionEntryProvenance::LegacyUnknown,
HistoryEntry::new(item, metadata)
}) })
.collect(); .collect();
Self::from_history_entries(segment_id, entries) Self::from_captured_entries(segment_id, entries)
} }
pub(crate) fn from_history_entries( pub(crate) fn from_history_entries(
segment_id: impl Into<String>, segment_id: impl Into<String>,
entries: Vec<HistoryEntry<SessionHistoryMetadata>>, entries: Vec<HistoryEntry<SessionHistoryMetadata>>,
) -> Self {
let entries = entries
.into_iter()
.map(|entry| CapturedHistoryEntry {
item: entry.item,
entry_id: entry.annotation.entry_id,
origin: public_provenance(&entry.annotation.origin),
})
.collect();
Self::from_captured_entries(segment_id, entries)
}
fn from_captured_entries(
segment_id: impl Into<String>,
entries: Vec<CapturedHistoryEntry>,
) -> Self { ) -> Self {
let segment_id = segment_id.into(); let segment_id = segment_id.into();
let entries = Arc::new(entries);
let mut overview = Vec::new(); let mut overview = Vec::new();
let mut index = Vec::new(); let mut index = Vec::new();
@@ -317,7 +338,7 @@ impl SessionCapture {
let entry_range = [idx as u64, idx as u64]; let entry_range = [idx as u64, idx as u64];
match item { match item {
Item::Message { role, content, .. } => { Item::Message { role, content, .. } => {
let Some(kind) = message_reference_kind(&entry.annotation.origin, role) else { let Some(kind) = message_reference_kind(&entry.origin, role) else {
continue; continue;
}; };
let text = content let text = content
@@ -327,10 +348,10 @@ impl SessionCapture {
.join(""); .join("");
let label = format!("{} message", kind.as_str()); let label = format!("{} message", kind.as_str());
let summary = truncate_chars(&text, 240); let summary = truncate_chars(&text, 240);
let id = SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id); let id = SessionEntryRef::from_history_entry_id(&entry.entry_id);
index.push(ReferenceEntry { index.push(ReferenceEntry {
id: id.clone(), id: id.clone(),
origin: entry.annotation.origin.clone(), origin: entry.origin.clone(),
entry_range, entry_range,
kind, kind,
tool_part: None, tool_part: None,
@@ -342,7 +363,7 @@ impl SessionCapture {
if matches!(kind, ReferenceKind::User | ReferenceKind::Assistant) { if matches!(kind, ReferenceKind::User | ReferenceKind::Assistant) {
overview.push(OverviewItem { overview.push(OverviewItem {
id: id.clone(), id: id.clone(),
origin: entry.annotation.origin.clone(), origin: entry.origin.clone(),
entry_range, entry_range,
kind, kind,
label, label,
@@ -356,8 +377,8 @@ impl SessionCapture {
} => { } => {
let text = format!("{name}\n{arguments}"); let text = format!("{name}\n{arguments}");
index.push(ReferenceEntry { index.push(ReferenceEntry {
id: SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id), id: SessionEntryRef::from_history_entry_id(&entry.entry_id),
origin: entry.annotation.origin.clone(), origin: entry.origin.clone(),
entry_range, entry_range,
kind: ReferenceKind::Tool, kind: ReferenceKind::Tool,
tool_part: Some(ToolPart::Input), tool_part: Some(ToolPart::Input),
@@ -383,8 +404,8 @@ impl SessionCapture {
content.as_deref().unwrap_or_default(), content.as_deref().unwrap_or_default(),
); );
index.push(ReferenceEntry { index.push(ReferenceEntry {
id: SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id), id: SessionEntryRef::from_history_entry_id(&entry.entry_id),
origin: entry.annotation.origin.clone(), origin: entry.origin.clone(),
entry_range, entry_range,
kind: ReferenceKind::Tool, kind: ReferenceKind::Tool,
tool_part: Some(ToolPart::Output), tool_part: Some(ToolPart::Output),
@@ -424,7 +445,7 @@ impl SessionCapture {
Self { Self {
segment_id, segment_id,
entries, entries: Arc::new(entries),
overview, overview,
index, index,
} }
@@ -607,65 +628,41 @@ impl SessionCapture {
} }
} }
fn public_snapshot_metadata( fn public_provenance(origin: &WorkerHistoryProvenance) -> SessionEntryProvenance {
entry_id: String, match origin {
provenance: SessionEntryProvenance, WorkerHistoryProvenance::HumanInput { .. } => SessionEntryProvenance::HumanInput,
) -> SessionHistoryMetadata { WorkerHistoryProvenance::WorkerInput { .. } => SessionEntryProvenance::WorkerInput,
let worker = session_store::LoggedWorkerSubject { WorkerHistoryProvenance::FlowInstruction { .. } => SessionEntryProvenance::FlowInstruction,
workspace_id: None, WorkerHistoryProvenance::BackendInstruction { .. } => {
runtime_id: None, SessionEntryProvenance::BackendInstruction
worker_id: "public-session-snapshot".to_owned(),
};
let origin = match provenance {
SessionEntryProvenance::HumanInput => WorkerHistoryProvenance::HumanInput {
account_id: "public-session-snapshot".to_owned(),
},
SessionEntryProvenance::WorkerInput => WorkerHistoryProvenance::WorkerInput {
actor: worker.clone(),
},
SessionEntryProvenance::FlowInstruction => WorkerHistoryProvenance::FlowInstruction {
selector: "public-session-snapshot".to_owned(),
definition_id: "public-session-snapshot".to_owned(),
definition_revision: 0,
instance_id: "public-session-snapshot".to_owned(),
state_id: "public-session-snapshot".to_owned(),
},
SessionEntryProvenance::BackendInstruction => {
WorkerHistoryProvenance::BackendInstruction { operation_id: None }
} }
SessionEntryProvenance::ModelOutput => WorkerHistoryProvenance::ModelOutput { WorkerHistoryProvenance::ModelOutput { .. } => SessionEntryProvenance::ModelOutput,
worker: worker.clone(), WorkerHistoryProvenance::ToolOutput { .. } => SessionEntryProvenance::ToolOutput,
}, WorkerHistoryProvenance::DerivedSummary => SessionEntryProvenance::DerivedSummary,
SessionEntryProvenance::ToolOutput => WorkerHistoryProvenance::ToolOutput { worker }, WorkerHistoryProvenance::LegacyUnknown => SessionEntryProvenance::LegacyUnknown,
SessionEntryProvenance::DerivedSummary => WorkerHistoryProvenance::DerivedSummary,
SessionEntryProvenance::LegacyUnknown => WorkerHistoryProvenance::LegacyUnknown,
};
SessionHistoryMetadata {
entry_id: session_store::LoggedSessionHistoryEntryId(entry_id),
origin,
derivation: None,
} }
} }
fn message_reference_kind( fn message_reference_kind(
origin: &WorkerHistoryProvenance, origin: &SessionEntryProvenance,
provider_role: &Role, provider_role: &Role,
) -> Option<ReferenceKind> { ) -> Option<ReferenceKind> {
match origin { match origin {
WorkerHistoryProvenance::HumanInput { .. } SessionEntryProvenance::HumanInput | SessionEntryProvenance::WorkerInput => {
| WorkerHistoryProvenance::WorkerInput { .. } => Some(ReferenceKind::User), Some(ReferenceKind::User)
WorkerHistoryProvenance::ModelOutput { .. } => Some(ReferenceKind::Assistant), }
WorkerHistoryProvenance::ToolOutput { .. } => Some(ReferenceKind::Tool), SessionEntryProvenance::ModelOutput => Some(ReferenceKind::Assistant),
WorkerHistoryProvenance::LegacyUnknown => match provider_role { SessionEntryProvenance::ToolOutput => Some(ReferenceKind::Tool),
SessionEntryProvenance::LegacyUnknown => match provider_role {
Role::User => Some(ReferenceKind::User), Role::User => Some(ReferenceKind::User),
Role::Assistant => Some(ReferenceKind::Assistant), Role::Assistant => Some(ReferenceKind::Assistant),
Role::System => None, Role::System => None,
}, },
// Flow/backend/system content remains out of the observation surface // Flow/backend/system content remains out of the observation surface
// even when represented with a provider user/system role. // even when represented with a provider user/system role.
WorkerHistoryProvenance::FlowInstruction { .. } SessionEntryProvenance::FlowInstruction
| WorkerHistoryProvenance::BackendInstruction { .. } | SessionEntryProvenance::BackendInstruction
| WorkerHistoryProvenance::DerivedSummary => None, | SessionEntryProvenance::DerivedSummary => None,
} }
} }
@@ -762,13 +759,13 @@ mod tests {
assert_eq!(overview.len(), 1); assert_eq!(overview.len(), 1);
assert!(matches!( assert!(matches!(
overview[0].origin, overview[0].origin,
WorkerHistoryProvenance::HumanInput { .. } SessionEntryProvenance::HumanInput
)); ));
let evidence = capture.evidence_for(overview[0].id.as_str()).unwrap(); let evidence = capture.evidence_for(overview[0].id.as_str()).unwrap();
assert!(evidence.excerpt.ends_with("remember my preference")); assert!(evidence.excerpt.ends_with("remember my preference"));
assert!(matches!( assert!(matches!(
evidence.origin, evidence.origin,
WorkerHistoryProvenance::HumanInput { .. } SessionEntryProvenance::HumanInput
)); ));
} }
crates/worker/src/spawn/comm_tools.rs
+1
View File
@@ -284,6 +284,7 @@ mod tests {
.enumerate() .enumerate()
.map(|(index, value)| protocol::SessionSnapshotEntry { .map(|(index, value)| protocol::SessionSnapshotEntry {
entry_id: format!("test-{index}"), entry_id: format!("test-{index}"),
timestamp: index as u64,
provenance: protocol::SessionEntryProvenance::LegacyUnknown, provenance: protocol::SessionEntryProvenance::LegacyUnknown,
derived_from: Vec::new(), derived_from: Vec::new(),
data: protocol::SessionSnapshotEntryData::RunError { data: protocol::SessionSnapshotEntryData::RunError {
web/workspace/src/lib/generated/protocol.ts
+5 -1
View File
@@ -93,7 +93,11 @@ export type SessionSnapshotEntry = {
* Stable identity from durable history metadata, or a deterministic * Stable identity from durable history metadata, or a deterministic
* identity derived from the legacy segment and log position. * identity derived from the legacy segment and log position.
*/ */
entry_id: string, provenance: SessionEntryProvenance, derived_from?: Array<string>, } & ({ "kind": "user_input", segments: Array<Segment>, } | { "kind": "message", role: SessionMessageRole, content: Array<SessionContentPart>, } | { "kind": "tool_call", call_id: string, name: string, arguments: string, } | { "kind": "tool_result", call_id: string, summary: string, content?: string | null, is_error: boolean, attachments?: Array<SessionToolAttachment>, } | { "kind": "system_item", item_kind: string, content: string, data?: unknown, } | { "kind": "run_error", message: string, }); entry_id: string,
/**
* Timestamp copied from the durable log record that commits this entry.
*/
timestamp: number, provenance: SessionEntryProvenance, derived_from?: Array<string>, } & ({ "kind": "user_input", segments: Array<Segment>, } | { "kind": "message", role: SessionMessageRole, content: Array<SessionContentPart>, } | { "kind": "tool_call", call_id: string, name: string, arguments: string, } | { "kind": "tool_result", call_id: string, summary: string, content?: string | null, is_error: boolean, attachments?: Array<SessionToolAttachment>, } | { "kind": "system_item", item_kind: string, content: string, data?: unknown, } | { "kind": "run_error", message: string, });
export type SessionSnapshot = { entries: Array<SessionSnapshotEntry>, }; export type SessionSnapshot = { entries: Array<SessionSnapshotEntry>, };
web/workspace/src/lib/workspace/console/model.test.ts
+1
View File
@@ -2108,6 +2108,7 @@ Deno.test("snapshot restores TaskStore state from system history", () => {
event.data.session = { event.data.session = {
entries: [{ entries: [{
entry_id: "task-reminder-1", entry_id: "task-reminder-1",
timestamp: 1,
provenance: "backend_instruction", provenance: "backend_instruction",
kind: "system_item", kind: "system_item",
item_kind: "task_reminder", item_kind: "task_reminder",