From 88be87e03e6c9705aa48410fc0192b3d2b85ec23 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 30 Aug 2026 00:41:51 +0900 Subject: [PATCH 1/4] feat(session): expose canonical public snapshots --- crates/protocol/src/lib.rs | 163 +++++-- crates/protocol/src/typescript.rs | 11 +- crates/session-store/src/history.rs | 24 +- crates/session-store/src/lib.rs | 1 + crates/session-store/src/public_snapshot.rs | 402 ++++++++++++++++++ .../session-store/src/worker_session_store.rs | 396 ++++++++++++++++- crates/tui/src/app.rs | 283 ++++++------ crates/tui/src/console/mod.rs | 16 +- crates/tui/src/dashboard/tests.rs | 8 +- crates/tui/src/worker_list.rs | 2 +- crates/worker-runtime/src/runtime.rs | 24 +- crates/worker/src/controller.rs | 41 +- crates/worker/src/discovery.rs | 24 +- .../src/feature/builtin/worker_observation.rs | 90 ++-- crates/worker/src/internal_worker.rs | 4 +- crates/worker/src/ipc/protocol_session.rs | 5 +- crates/worker/src/session_capture.rs | 104 +++++ crates/worker/src/spawn/comm_tools.rs | 15 +- crates/worker/src/spawn/registry.rs | 8 +- crates/worker/src/spawn/tool.rs | 12 +- crates/worker/tests/controller_test.rs | 49 +-- .../tests/system_prompt_template_test.rs | 4 +- crates/workspace-server/src/server.rs | 6 +- web/workspace/src/lib/generated/protocol.ts | 30 +- .../src/lib/workspace/console/model.test.ts | 127 ++++-- .../src/lib/workspace/console/model.ts | 76 ++-- 26 files changed, 1512 insertions(+), 413 deletions(-) create mode 100644 crates/session-store/src/public_snapshot.rs diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index bfc1a49a..339d1d67 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -340,8 +340,7 @@ pub struct InternalWorkerRef { pub struct InternalWorkerSnapshot { pub worker: InternalWorkerRef, pub revision: u64, - #[cfg_attr(feature = "typescript", ts(type = "Array"))] - pub entries: Vec, + pub session: SessionSnapshot, #[serde(default)] pub status: WorkerStatus, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -364,6 +363,106 @@ pub enum ToolResultDisposition { OutcomeUnknown, } +/// Canonical, storage-independent projection of committed session history. +/// +/// Worker protocols expose this DTO instead of append-log records. New +/// storage variants can therefore be added without teaching every client how +/// to replay the durable log format. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +pub struct SessionSnapshot { + pub entries: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(rename_all = "snake_case")] +pub enum SessionEntryProvenance { + HumanInput, + WorkerInput, + FlowInstruction, + BackendInstruction, + ModelOutput, + ToolOutput, + DerivedSummary, + LegacyUnknown, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +pub struct SessionSnapshotEntry { + /// Stable identity from durable history metadata, or a deterministic + /// identity derived from the legacy segment and log position. + pub entry_id: String, + pub provenance: SessionEntryProvenance, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub derived_from: Vec, + #[serde(flatten)] + pub data: SessionSnapshotEntryData, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum SessionSnapshotEntryData { + UserInput { + segments: Vec, + }, + Message { + role: SessionMessageRole, + content: Vec, + }, + ToolCall { + call_id: String, + name: String, + arguments: String, + }, + ToolResult { + call_id: String, + summary: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + content: Option, + is_error: bool, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + attachments: Vec, + }, + SystemItem { + item_kind: String, + content: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "typescript", ts(type = "unknown"))] + data: Option, + }, + RunError { + message: String, + }, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(rename_all = "snake_case")] +pub enum SessionMessageRole { + User, + Assistant, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum SessionContentPart { + Text { text: String }, + Refusal { refusal: String }, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +pub struct SessionToolAttachment { + pub media_type: String, + /// Base64-encoded durable attachment body. Public snapshots preserve the + /// committed multimodal value instead of replacing it with placeholder text. + pub data_base64: String, +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "typescript", derive(ts_rs::TS))] #[serde(tag = "event", content = "data", rename_all = "snake_case")] @@ -555,8 +654,7 @@ pub enum Event { /// role-specific entry events (`SegmentRotated` / `SystemItem`) — /// there is no generic "every committed entry" broadcast. Snapshot { - #[cfg_attr(feature = "typescript", ts(type = "Array"))] - entries: Vec, + session: SessionSnapshot, greeting: Greeting, #[serde(default)] status: WorkerStatus, @@ -589,14 +687,10 @@ pub enum Event { /// Server-side segment log rotated to a fresh `SegmentStart`. /// /// Fires on compaction and on auto-fork when the store head drifts - /// from the live writer's cached head. Clients drop their derived - /// view and reseed from `entry.history` exactly the way they would - /// from a connect-time `Snapshot`. - /// - /// Payload is the JSON form of `session_store::LogEntry::SegmentStart`. + /// A compaction/fork has replaced the authoritative segment. Clients drop + /// their derived view and reseed from the canonical committed snapshot. SegmentRotated { - #[cfg_attr(feature = "typescript", ts(type = "unknown"))] - entry: serde_json::Value, + session: SessionSnapshot, }, /// Current Worker controller status. Broadcast on every controller-level /// transition and included in `History` snapshots for late attach. @@ -623,11 +717,10 @@ pub enum Event { head_entries: usize, targets: Vec, }, - /// A rewind has truncated the authoritative session. `entries` is the - /// retained session-log prefix clients should use to reseed display state. + /// A rewind has truncated the authoritative session. `session` is the + /// retained canonical snapshot clients should use to reseed display state. RewindApplied { - #[cfg_attr(feature = "typescript", ts(type = "Array"))] - entries: Vec, + session: SessionSnapshot, input: Vec, summary: RewindSummary, }, @@ -1440,7 +1533,16 @@ mod tests { #[test] fn event_snapshot_format() { let event = Event::Snapshot { - entries: vec![serde_json::json!({"kind": "user_input", "ts": 1, "segments": []})], + session: SessionSnapshot { + entries: vec![SessionSnapshotEntry { + entry_id: "entry-1".into(), + provenance: SessionEntryProvenance::HumanInput, + derived_from: Vec::new(), + data: SessionSnapshotEntryData::UserInput { + segments: Vec::new(), + }, + }], + }, greeting: Greeting { worker_name: "test".into(), cwd: "/tmp".into(), @@ -1458,8 +1560,11 @@ mod tests { let json = serde_json::to_string(&event).unwrap(); let parsed: serde_json::Value = serde_json::from_str(&json).unwrap(); assert_eq!(parsed["event"], "snapshot"); - assert!(parsed["data"]["entries"].is_array()); - assert_eq!(parsed["data"]["entries"][0]["kind"], "user_input"); + assert!(parsed["data"]["session"]["entries"].is_array()); + assert_eq!( + parsed["data"]["session"]["entries"][0]["kind"], + "user_input" + ); assert_eq!(parsed["data"]["greeting"]["worker_name"], "test"); assert_eq!(parsed["data"]["greeting"]["tools"][0], "Read"); assert_eq!(parsed["data"]["greeting"]["context_window"], 200_000); @@ -1469,7 +1574,7 @@ mod tests { #[test] fn event_snapshot_in_flight_roundtrip_and_default() { - let inbound = r#"{"event":"snapshot","data":{"entries":[],"greeting":{"worker_name":"test","cwd":"/tmp","provider":"p","model":"m","scope_summary":"s","tools":[]},"status":"running"}}"#; + let inbound = r#"{"event":"snapshot","data":{"session":{"entries":[]},"greeting":{"worker_name":"test","cwd":"/tmp","provider":"p","model":"m","scope_summary":"s","tools":[]},"status":"running"}}"#; let decoded: Event = serde_json::from_str(inbound).unwrap(); match decoded { Event::Snapshot { in_flight, .. } => assert!(in_flight.is_empty()), @@ -1477,7 +1582,9 @@ mod tests { } let event = Event::Snapshot { - entries: Vec::new(), + session: SessionSnapshot { + entries: Vec::new(), + }, greeting: Greeting { worker_name: "test".into(), cwd: "/tmp".into(), @@ -1543,15 +1650,17 @@ mod tests { #[test] fn event_segment_rotated_roundtrip() { let event = Event::SegmentRotated { - entry: serde_json::json!({"kind": "segment_start", "ts": 1, "history": []}), + session: SessionSnapshot { + entries: Vec::new(), + }, }; let json = serde_json::to_string(&event).unwrap(); let parsed: serde_json::Value = serde_json::from_str(&json).unwrap(); assert_eq!(parsed["event"], "segment_rotated"); - assert_eq!(parsed["data"]["entry"]["kind"], "segment_start"); + assert!(parsed["data"]["session"]["entries"].is_array()); let decoded: Event = serde_json::from_str(&json).unwrap(); match decoded { - Event::SegmentRotated { entry } => assert_eq!(entry["kind"], "segment_start"), + Event::SegmentRotated { session } => assert!(session.entries.is_empty()), other => panic!("expected SegmentRotated, got {other:?}"), } } @@ -1627,8 +1736,8 @@ mod tests { } #[test] - fn event_snapshot_legacy_without_status_defaults_to_idle() { - let json = r#"{"event":"snapshot","data":{"entries":[],"greeting":{"worker_name":"test","cwd":"/tmp","provider":"anthropic","model":"claude","scope_summary":"","tools":[]}}}"#; + fn event_snapshot_without_status_defaults_to_idle() { + let json = r#"{"event":"snapshot","data":{"session":{"entries":[]},"greeting":{"worker_name":"test","cwd":"/tmp","provider":"anthropic","model":"claude","scope_summary":"","tools":[]}}}"#; let decoded: Event = serde_json::from_str(json).unwrap(); match decoded { Event::Snapshot { @@ -2039,11 +2148,11 @@ mod tests { } #[test] - fn legacy_snapshot_defaults_internal_workers_to_empty() { + fn snapshot_defaults_internal_workers_to_empty() { let snapshot: Event = serde_json::from_value(serde_json::json!({ "event": "snapshot", "data": { - "entries": [], + "session": { "entries": [] }, "greeting": { "worker_name": "parent", "cwd": ".", diff --git a/crates/protocol/src/typescript.rs b/crates/protocol/src/typescript.rs index 2ac5446c..bf75e324 100644 --- a/crates/protocol/src/typescript.rs +++ b/crates/protocol/src/typescript.rs @@ -8,7 +8,9 @@ use crate::{ CompletionKind, ErrorCode, Event, Greeting, InFlightBlock, InFlightSnapshot, InFlightToolCallState, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnapshot, InvokeKind, MemoryWorkerEvent, Method, Permission, RewindSummary, RewindTarget, RewindTargetId, - RunResult, ScopeRule, Segment, ToolResultDisposition, TurnResult, WorkerEvent, WorkerStatus, + RunResult, ScopeRule, Segment, SessionContentPart, SessionEntryProvenance, SessionMessageRole, + SessionSnapshot, SessionSnapshotEntry, SessionSnapshotEntryData, SessionToolAttachment, + ToolResultDisposition, TurnResult, WorkerEvent, WorkerStatus, subscription::{ EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame, SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest, @@ -63,6 +65,13 @@ pub fn generated_protocol_types() -> String { push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); diff --git a/crates/session-store/src/history.rs b/crates/session-store/src/history.rs index 3903213c..9558f343 100644 --- a/crates/session-store/src/history.rs +++ b/crates/session-store/src/history.rs @@ -1,5 +1,6 @@ //! Serializable history entries with restore-authoritative logical identity and origin. +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use serde::{Deserialize, Serialize}; use crate::{LoggedItem, SessionId}; @@ -175,6 +176,25 @@ pub fn legacy_segment_history( session_id: SessionId, items: impl IntoIterator, ) -> Vec { - let _ = session_id; - items.into_iter().map(legacy_logged_history).collect() + items + .into_iter() + .enumerate() + .map(|(index, item)| LoggedHistoryEntry { + item, + metadata: LoggedSessionHistoryMetadata { + // Legacy logs have no persisted entry id. Derive one solely from + // durable segment content rather than minting a new random value + // on every restore/read. The explicit LegacyUnknown origin keeps + // this compatibility identity from becoming trust authority. + entry_id: { + let mut identity = Vec::with_capacity(24); + identity.extend_from_slice(session_id.as_bytes()); + identity.extend_from_slice(&(index as u64).to_be_bytes()); + LoggedSessionHistoryEntryId(format!("l-{}", URL_SAFE_NO_PAD.encode(identity))) + }, + origin: LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + }, + }) + .collect() } diff --git a/crates/session-store/src/lib.rs b/crates/session-store/src/lib.rs index 9fe1fad7..07bc0bc9 100644 --- a/crates/session-store/src/lib.rs +++ b/crates/session-store/src/lib.rs @@ -34,6 +34,7 @@ pub mod event_trace; pub mod fs_store; pub mod history; pub mod logged_item; +pub mod public_snapshot; pub mod segment; pub mod segment_log; pub mod store; diff --git a/crates/session-store/src/public_snapshot.rs b/crates/session-store/src/public_snapshot.rs new file mode 100644 index 00000000..d25feb99 --- /dev/null +++ b/crates/session-store/src/public_snapshot.rs @@ -0,0 +1,402 @@ +use base64::{ + Engine as _, + engine::general_purpose::{STANDARD as BASE64, URL_SAFE_NO_PAD}, +}; +use protocol::{ + Segment, SessionContentPart, SessionEntryProvenance, SessionMessageRole, SessionSnapshot, + SessionSnapshotEntry, SessionSnapshotEntryData, SessionToolAttachment, +}; + +use crate::{ + LogEntry, LoggedContentPart, LoggedHistoryEntry, LoggedItem, LoggedRole, + LoggedSessionHistoryOrigin, SessionId, SystemItem, +}; + +/// Project a complete current-segment log. A valid segment always starts with +/// one of the two SegmentStart records; malformed partial input uses the nil +/// session only to keep the public failure projection deterministic. +pub fn project_current_session_snapshot(log: &[LogEntry]) -> SessionSnapshot { + let session_id = log.iter().find_map(|entry| match entry { + LogEntry::SegmentStart { session_id, .. } + | LogEntry::AnnotatedSegmentStart { session_id, .. } => Some(*session_id), + _ => None, + }); + project_session_snapshot(session_id.unwrap_or_else(SessionId::nil), log) +} + +/// Project the current durable segment into the only public session-history +/// representation. Append-log records remain an internal persistence format. +pub fn project_session_snapshot(session_id: SessionId, log: &[LogEntry]) -> SessionSnapshot { + let mut session_key = session_id; + let mut entries = Vec::new(); + + for (log_index, record) in log.iter().enumerate() { + match record { + LogEntry::SegmentStart { + session_id, + history, + .. + } => { + session_key = *session_id; + 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)); + } + } + } + LogEntry::AnnotatedSegmentStart { + session_id, + history, + .. + } => { + session_key = *session_id; + entries.clear(); + extend_history(&mut entries, history, None); + } + LogEntry::UserInput { segments, .. } => entries.push(legacy_entry( + &session_key, + log_index, + 0, + SessionSnapshotEntryData::UserInput { + segments: segments.clone(), + }, + )), + LogEntry::AnnotatedUserInput { + segments, history, .. + } => extend_history(&mut entries, history, Some(segments)), + LogEntry::AssistantItem { item, .. } | LogEntry::ToolResult { item, .. } => { + if let Some(data) = project_item(item) { + entries.push(legacy_entry(&session_key, log_index, 0, data)); + } + } + LogEntry::AnnotatedAssistantItem { entry, .. } + | LogEntry::AnnotatedToolResult { entry, .. } => { + if let Some(data) = project_item(&entry.item) { + entries.push(history_entry(entry, data)); + } + } + LogEntry::SystemItem { item, .. } => entries.push(system_entry( + item, + legacy_entry_id(&session_key, log_index, 0), + SessionEntryProvenance::LegacyUnknown, + Vec::new(), + )), + LogEntry::AnnotatedSystemItem { entry, .. } => entries.push(system_entry( + &entry.item, + entry.metadata.entry_id.0.clone(), + provenance(&entry.metadata.origin), + derivation_ids(entry), + )), + LogEntry::RunErrored { message, .. } => entries.push(legacy_entry( + &session_key, + log_index, + 0, + SessionSnapshotEntryData::RunError { + message: message.clone(), + }, + )), + // Run checkpoints, configuration, usage, and extension state are + // controller/storage authority rather than committed conversation. + LogEntry::Invoke { .. } + | LogEntry::TurnEnd { .. } + | LogEntry::RunCompleted { .. } + | LogEntry::ActiveRunCheckpoint { .. } + | LogEntry::PausedTurnAbandoned { .. } + | LogEntry::ConfigChanged { .. } + | LogEntry::LlmUsage { .. } + | LogEntry::Extension { .. } => {} + } + } + + SessionSnapshot { entries } +} + +fn extend_history( + output: &mut Vec, + history: &[LoggedHistoryEntry], + input_segments: Option<&Vec>, +) { + let mut attached_segments = false; + for entry in history { + let data = if !attached_segments + && matches!( + entry.metadata.origin, + LoggedSessionHistoryOrigin::HumanInput { .. } + ) + && input_segments.is_some() + { + attached_segments = true; + SessionSnapshotEntryData::UserInput { + segments: input_segments.cloned().unwrap_or_default(), + } + } else { + let Some(data) = project_item(&entry.item) else { + continue; + }; + data + }; + output.push(history_entry(entry, data)); + } +} + +fn history_entry( + entry: &LoggedHistoryEntry, + data: SessionSnapshotEntryData, +) -> SessionSnapshotEntry { + SessionSnapshotEntry { + entry_id: entry.metadata.entry_id.0.clone(), + provenance: provenance(&entry.metadata.origin), + derived_from: entry + .metadata + .derivation + .as_ref() + .map(|derivation| { + derivation + .sources + .iter() + .map(|source| source.0.clone()) + .collect() + }) + .unwrap_or_default(), + data, + } +} + +fn derivation_ids(entry: &crate::LoggedSystemHistoryEntry) -> Vec { + entry + .metadata + .derivation + .as_ref() + .map(|derivation| { + derivation + .sources + .iter() + .map(|source| source.0.clone()) + .collect() + }) + .unwrap_or_default() +} + +fn legacy_entry( + session_key: &SessionId, + log_index: usize, + item_index: usize, + data: SessionSnapshotEntryData, +) -> SessionSnapshotEntry { + SessionSnapshotEntry { + entry_id: legacy_entry_id(session_key, log_index, item_index), + provenance: SessionEntryProvenance::LegacyUnknown, + derived_from: Vec::new(), + data, + } +} + +fn legacy_entry_id(session_key: &SessionId, log_index: usize, item_index: usize) -> String { + let mut identity = Vec::with_capacity(32); + identity.extend_from_slice(session_key.as_bytes()); + identity.extend_from_slice(&(log_index as u64).to_be_bytes()); + identity.extend_from_slice(&(item_index as u64).to_be_bytes()); + format!("l-{}", URL_SAFE_NO_PAD.encode(identity)) +} + +fn provenance(origin: &LoggedSessionHistoryOrigin) -> SessionEntryProvenance { + match origin { + LoggedSessionHistoryOrigin::HumanInput { .. } => SessionEntryProvenance::HumanInput, + LoggedSessionHistoryOrigin::WorkerInput { .. } => SessionEntryProvenance::WorkerInput, + LoggedSessionHistoryOrigin::FlowInstruction { .. } => { + SessionEntryProvenance::FlowInstruction + } + LoggedSessionHistoryOrigin::BackendInstruction { .. } => { + SessionEntryProvenance::BackendInstruction + } + LoggedSessionHistoryOrigin::ModelOutput { .. } => SessionEntryProvenance::ModelOutput, + LoggedSessionHistoryOrigin::ToolOutput { .. } => SessionEntryProvenance::ToolOutput, + LoggedSessionHistoryOrigin::DerivedSummary => SessionEntryProvenance::DerivedSummary, + LoggedSessionHistoryOrigin::LegacyUnknown => SessionEntryProvenance::LegacyUnknown, + } +} + +fn project_item(item: &LoggedItem) -> Option { + match item { + LoggedItem::Message { role, content } => { + let role = match role { + LoggedRole::User => SessionMessageRole::User, + LoggedRole::Assistant => SessionMessageRole::Assistant, + // System prompts and instruction history never cross the public + // snapshot boundary. Typed SystemItems have separate records. + LoggedRole::System => return None, + }; + Some(SessionSnapshotEntryData::Message { + role, + content: content + .iter() + .map(|part| match part { + LoggedContentPart::Text { text } => { + SessionContentPart::Text { text: text.clone() } + } + LoggedContentPart::Refusal { refusal } => SessionContentPart::Refusal { + refusal: refusal.clone(), + }, + }) + .collect(), + }) + } + LoggedItem::ToolCall { + call_id, + name, + arguments, + } => Some(SessionSnapshotEntryData::ToolCall { + call_id: call_id.clone(), + name: name.clone(), + arguments: arguments.clone(), + }), + LoggedItem::ToolResult { + call_id, + summary, + content, + is_error, + attachments, + .. + } => Some(SessionSnapshotEntryData::ToolResult { + call_id: call_id.clone(), + summary: summary.clone(), + content: content.clone(), + is_error: *is_error, + attachments: attachments + .iter() + .map(|attachment| match attachment { + crate::logged_item::LoggedAttachment::Image { mime_type, data } => { + SessionToolAttachment { + media_type: mime_type.clone(), + data_base64: BASE64.encode(data), + } + } + }) + .collect(), + }), + // Hidden model reasoning is never observable. + LoggedItem::Reasoning { .. } => None, + } +} + +fn system_entry( + item: &SystemItem, + entry_id: String, + provenance: SessionEntryProvenance, + derived_from: Vec, +) -> SessionSnapshotEntry { + let mut data = serde_json::to_value(item).ok(); + if let Some(serde_json::Value::Object(object)) = data.as_mut() { + object.remove("prompt_provenance"); + } + let item_kind = data + .as_ref() + .and_then(|value| value.get("kind")) + .and_then(serde_json::Value::as_str) + .unwrap_or("system_item") + .to_owned(); + SessionSnapshotEntry { + entry_id, + provenance, + derived_from, + data: SessionSnapshotEntryData::SystemItem { + item_kind, + content: item.history_text(), + data, + }, + } +} + +#[cfg(test)] +mod tests { + use agen::llm_client::RequestConfig; + + use super::*; + use crate::{LoggedSessionHistoryEntryId, LoggedSessionHistoryMetadata, LoggedWorkerSubject}; + + #[test] + fn legacy_projection_is_stable_and_hides_reasoning_and_system_prompts() { + let session_id = crate::new_session_id(); + let log = vec![LogEntry::SegmentStart { + ts: 1, + session_id, + system_prompt: None, + config: RequestConfig::default(), + history: vec![ + LoggedItem::Message { + role: LoggedRole::System, + content: vec![LoggedContentPart::Text { + text: "secret prompt".into(), + }], + }, + LoggedItem::Reasoning { + text: "secret reasoning".into(), + summary: Vec::new(), + encrypted_content: None, + signature: None, + }, + LoggedItem::Message { + role: LoggedRole::Assistant, + content: vec![LoggedContentPart::Text { + text: "visible".into(), + }], + }, + ], + forked_from: None, + compacted_from: None, + }]; + + let first = project_session_snapshot(session_id, &log); + let second = project_session_snapshot(session_id, &log); + assert_eq!(first, second); + assert_eq!(first.entries.len(), 1); + assert_eq!( + first.entries[0].provenance, + SessionEntryProvenance::LegacyUnknown + ); + let json = serde_json::to_string(&first).unwrap(); + assert!(!json.contains("secret prompt")); + assert!(!json.contains("secret reasoning")); + assert!(json.contains("visible")); + } + + #[test] + fn annotated_projection_preserves_identity_and_provenance() { + let session_id = crate::new_session_id(); + let metadata = LoggedSessionHistoryMetadata { + entry_id: LoggedSessionHistoryEntryId::new(), + origin: LoggedSessionHistoryOrigin::ModelOutput { + worker: LoggedWorkerSubject { + workspace_id: None, + runtime_id: None, + worker_id: "worker".into(), + }, + }, + derivation: None, + }; + let expected_id = metadata.entry_id.0.clone(); + let log = vec![LogEntry::AnnotatedSegmentStart { + ts: 1, + session_id, + system_prompt: None, + config: RequestConfig::default(), + history: vec![LoggedHistoryEntry { + item: LoggedItem::Message { + role: LoggedRole::Assistant, + content: vec![LoggedContentPart::Text { text: "ok".into() }], + }, + metadata, + }], + forked_from: None, + compacted_from: None, + }]; + + let snapshot = project_session_snapshot(session_id, &log); + assert_eq!(snapshot.entries[0].entry_id, expected_id); + assert_eq!( + snapshot.entries[0].provenance, + SessionEntryProvenance::ModelOutput + ); + } +} diff --git a/crates/session-store/src/worker_session_store.rs b/crates/session-store/src/worker_session_store.rs index 5dd063d8..ddce9e1b 100644 --- a/crates/session-store/src/worker_session_store.rs +++ b/crates/session-store/src/worker_session_store.rs @@ -12,7 +12,11 @@ use crate::event_trace::TraceEntry; use crate::segment_log::LogEntry; use crate::store::{Store, StoreError}; -use crate::{SegmentId, SessionId}; +use crate::{ + LoggedHistoryEntry, LoggedItem, LoggedSessionHistoryEntryId, LoggedSessionHistoryMetadata, + LoggedSessionHistoryOrigin, LoggedSystemHistoryEntry, SegmentId, SessionId, +}; +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use serde::{Deserialize, Serialize}; use std::fs::{self, File, OpenOptions}; use std::io::{Read, Seek, SeekFrom, Write}; @@ -20,7 +24,8 @@ use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex}; use std::time::SystemTime; -const SESSION_SCHEMA_VERSION: u32 = 2; +const SESSION_SCHEMA_VERSION: u32 = 3; +const PREVIOUS_SESSION_SCHEMA_VERSION: u32 = 2; const LEGACY_SESSION_SCHEMA_VERSION: u32 = 1; const SESSION_FILE: &str = "session.json"; const SEGMENTS_DIR: &str = "segments"; @@ -47,9 +52,11 @@ impl WorkerSessionStore { Ok(bytes) => { let mut manifest: SessionManifest = serde_json::from_slice(&bytes)?; match manifest.schema_version { - SESSION_SCHEMA_VERSION => {} - LEGACY_SESSION_SCHEMA_VERSION => { - validate_legacy_segment_logs(&root)?; + SESSION_SCHEMA_VERSION => { + validate_canonical_segment_logs(&root)?; + } + PREVIOUS_SESSION_SCHEMA_VERSION | LEGACY_SESSION_SCHEMA_VERSION => { + migrate_segment_logs_to_v3(&root, manifest.session_id)?; manifest.schema_version = SESSION_SCHEMA_VERSION; atomic_write_json(&root.join(SESSION_FILE), &manifest)?; } @@ -144,6 +151,48 @@ impl WorkerSessionStore { .join(format!("{segment_id}.trace.jsonl")) } + fn append_log_entry( + &self, + path: &Path, + session_id: SessionId, + segment_id: SegmentId, + entry: &LogEntry, + ) -> Result<(), StoreError> { + let _guard = self + .append_lock + .lock() + .map_err(|_| std::io::Error::other("Worker Session append lock was poisoned"))?; + let mut file = OpenOptions::new() + .create(true) + .read(true) + .write(true) + .append(true) + .open(path)?; + let committed_len = truncate_uncommitted_tail(&mut file)?; + file.seek(SeekFrom::Start(0))?; + let mut existing = Vec::new(); + file.read_to_end(&mut existing)?; + let line_index = parse_jsonl::(&existing)?.len(); + let entry = canonicalize_log_entry(session_id, segment_id, line_index, entry.clone()); + let line = serde_json::to_string(&entry)?; + let mut record = Vec::with_capacity(line.len() + 1); + record.extend_from_slice(line.as_bytes()); + record.push(b'\n'); + if let Err(write_error) = file.write_all(&record) { + return match file.set_len(committed_len) { + Ok(()) => Err(write_error.into()), + Err(rollback_error) => Err(std::io::Error::new( + rollback_error.kind(), + format!( + "session append failed ({write_error}) and rollback failed: {rollback_error}" + ), + ) + .into()), + }; + } + Ok(()) + } + fn append_line(&self, path: &Path, line: &str) -> Result<(), StoreError> { let _guard = self .append_lock @@ -183,7 +232,7 @@ impl Store for WorkerSessionStore { entry: &LogEntry, ) -> Result<(), StoreError> { self.ensure_session(session_id, true)?; - self.append_line(&self.log_path(segment_id), &serde_json::to_string(entry)?) + self.append_log_entry(&self.log_path(segment_id), session_id, segment_id, entry) } fn read_all( @@ -236,8 +285,9 @@ impl Store for WorkerSessionStore { ) -> Result<(), StoreError> { self.ensure_session(session_id, true)?; let mut content = Vec::new(); - for entry in entries { - serde_json::to_writer(&mut content, entry)?; + for (line_index, entry) in entries.iter().enumerate() { + let entry = canonicalize_log_entry(session_id, segment_id, line_index, entry.clone()); + serde_json::to_writer(&mut content, &entry)?; content.push(b'\n'); } atomic_write_bytes(&self.log_path(segment_id), &content)?; @@ -286,37 +336,208 @@ impl Store for WorkerSessionStore { } } -fn validate_legacy_segment_logs(root: &Path) -> Result<(), StoreError> { +fn segment_log_paths(root: &Path) -> Result, StoreError> { let segments = root.join(SEGMENTS_DIR); if !segments.exists() { - return Ok(()); + return Ok(Vec::new()); } + let mut paths = Vec::new(); for entry in fs::read_dir(&segments)? { let entry = entry?; let path = entry.path(); + let metadata = fs::symlink_metadata(&path)?; let Some(name) = path.file_name().and_then(|name| name.to_str()) else { - continue; + return Err(StoreError::Corrupt { + line: 0, + message: format!("non-UTF-8 Worker Session segment path: {}", path.display()), + }); }; - if !name.ends_with(".jsonl") || name.ends_with(".trace.jsonl") { + if name.ends_with(".trace.jsonl") || name.starts_with('.') { continue; } - let contents = fs::read_to_string(&path)?; - for (line_index, line) in contents.lines().enumerate() { - if line.trim().is_empty() { - continue; - } - serde_json::from_str::(line).map_err(|error| StoreError::Corrupt { - line: line_index + 1, + if !name.ends_with(".jsonl") { + continue; + } + if !metadata.file_type().is_file() { + return Err(StoreError::Corrupt { + line: 0, message: format!( - "cannot migrate legacy Worker Session log {}: {error}", + "Worker Session segment is not a regular file: {}", path.display() ), - })?; + }); + } + let segment_id = + name.trim_end_matches(".jsonl") + .parse() + .map_err(|_| StoreError::Corrupt { + line: 0, + message: format!("invalid Worker Session segment name: {name}"), + })?; + paths.push((segment_id, path)); + } + paths.sort_by_key(|(segment_id, _)| *segment_id); + Ok(paths) +} + +fn migrate_segment_logs_to_v3(root: &Path, session_id: SessionId) -> Result<(), StoreError> { + for (segment_id, path) in segment_log_paths(root)? { + let source = fs::read(&path)?; + let entries: Vec = parse_jsonl(&source).map_err(|error| StoreError::Corrupt { + line: 0, + message: format!( + "cannot migrate Worker Session log {}: {error}", + path.display() + ), + })?; + let canonical = entries + .into_iter() + .enumerate() + .map(|(line_index, entry)| { + canonicalize_log_entry(session_id, segment_id, line_index, entry) + }) + .collect::>(); + validate_canonical_entries(&path, &canonical)?; + let mut output = Vec::new(); + for entry in canonical { + serde_json::to_writer(&mut output, &entry)?; + output.push(b'\n'); + } + + // 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 { + return Err(StoreError::Corrupt { + line: 0, + message: format!( + "Worker Session segment changed during migration: {}", + path.display() + ), + }); + } + atomic_write_bytes(&path, &output)?; + } + Ok(()) +} + +fn validate_canonical_segment_logs(root: &Path) -> Result<(), StoreError> { + for (_, path) in segment_log_paths(root)? { + let entries: Vec = parse_jsonl(&fs::read(&path)?)?; + validate_canonical_entries(&path, &entries)?; + } + Ok(()) +} + +fn validate_canonical_entries(path: &Path, entries: &[LogEntry]) -> Result<(), StoreError> { + for (line_index, entry) in entries.iter().enumerate() { + if matches!( + entry, + LogEntry::SegmentStart { .. } + | LogEntry::UserInput { .. } + | LogEntry::AssistantItem { .. } + | LogEntry::ToolResult { .. } + | LogEntry::SystemItem { .. } + ) { + return Err(StoreError::Corrupt { + line: line_index + 1, + message: format!( + "Worker Session schema v3 contains legacy history record in {}", + path.display() + ), + }); } } Ok(()) } +fn legacy_metadata( + _session_id: SessionId, + segment_id: SegmentId, + line_index: usize, + item_index: usize, +) -> LoggedSessionHistoryMetadata { + let mut identity = Vec::with_capacity(32); + identity.extend_from_slice(segment_id.as_bytes()); + identity.extend_from_slice(&(line_index as u64).to_be_bytes()); + identity.extend_from_slice(&(item_index as u64).to_be_bytes()); + LoggedSessionHistoryMetadata { + entry_id: LoggedSessionHistoryEntryId(format!("l-{}", URL_SAFE_NO_PAD.encode(identity))), + origin: LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + } +} + +fn canonicalize_log_entry( + session_id: SessionId, + segment_id: SegmentId, + line_index: usize, + entry: LogEntry, +) -> LogEntry { + match entry { + LogEntry::SegmentStart { + ts, + session_id, + system_prompt, + config, + history, + forked_from, + compacted_from, + } => LogEntry::AnnotatedSegmentStart { + ts, + session_id, + system_prompt, + config, + history: history + .into_iter() + .enumerate() + .map(|(item_index, item)| LoggedHistoryEntry { + item, + metadata: legacy_metadata(session_id, segment_id, line_index, item_index), + }) + .collect(), + forked_from, + compacted_from, + }, + LogEntry::UserInput { + ts, + segments, + extensions, + } => LogEntry::AnnotatedUserInput { + ts, + history: vec![LoggedHistoryEntry { + item: LoggedItem::from(agen::Item::user_message( + protocol::Segment::flatten_to_text(&segments), + )), + metadata: legacy_metadata(session_id, segment_id, line_index, 0), + }], + segments, + extensions, + }, + LogEntry::AssistantItem { ts, item } => LogEntry::AnnotatedAssistantItem { + ts, + entry: LoggedHistoryEntry { + item, + metadata: legacy_metadata(session_id, segment_id, line_index, 0), + }, + }, + LogEntry::ToolResult { ts, item } => LogEntry::AnnotatedToolResult { + ts, + entry: LoggedHistoryEntry { + item, + metadata: legacy_metadata(session_id, segment_id, line_index, 0), + }, + }, + LogEntry::SystemItem { ts, item } => LogEntry::AnnotatedSystemItem { + ts, + entry: LoggedSystemHistoryEntry { + item, + metadata: legacy_metadata(session_id, segment_id, line_index, 0), + }, + }, + canonical => canonical, + } +} + fn atomic_write_json(path: &Path, value: &T) -> Result<(), StoreError> { let mut bytes = serde_json::to_vec_pretty(value)?; bytes.push(b'\n'); @@ -445,7 +666,7 @@ mod tests { } #[test] - fn schema_v1_logs_are_validated_and_promoted_to_v2() { + fn schema_v1_logs_are_rewritten_and_promoted_to_v3() { let root = tempfile::tempdir().unwrap(); let session_id = new_session_id(); let segment_id = new_segment_id(); @@ -467,7 +688,7 @@ mod tests { } #[test] - fn schema_v1_migration_rejects_corrupt_log_before_manifest_update() { + fn schema_v1_migration_rejects_corrupt_log_before_v3_manifest_update() { let root = tempfile::tempdir().unwrap(); let session_id = new_session_id(); let manifest = SessionManifest { @@ -492,6 +713,135 @@ mod tests { assert_eq!(persisted.schema_version, LEGACY_SESSION_SCHEMA_VERSION); } + #[test] + fn schema_v2_migration_rewrites_legacy_records_with_stable_unknown_provenance() { + let root = tempfile::tempdir().unwrap(); + let session_id = new_session_id(); + let segment_id = new_segment_id(); + 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 source = vec![ + 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, + }, + LogEntry::UserInput { + ts: 2, + segments: vec![protocol::Segment::Text { + content: "hello".into(), + }], + extensions: Vec::new(), + }, + LogEntry::AssistantItem { + ts: 3, + item: LoggedItem::from(agen::Item::assistant_message("reply")), + }, + ]; + let path = root + .path() + .join(SEGMENTS_DIR) + .join(format!("{segment_id}.jsonl")); + let mut bytes = Vec::new(); + for entry in source { + serde_json::to_writer(&mut bytes, &entry).unwrap(); + bytes.push(b'\n'); + } + fs::write(&path, bytes).unwrap(); + + let store = WorkerSessionStore::new(root.path()).unwrap(); + let first = store.read_all(session_id, segment_id).unwrap(); + assert!(matches!(first[0], LogEntry::AnnotatedSegmentStart { .. })); + assert!(matches!(first[1], LogEntry::AnnotatedUserInput { .. })); + assert!(matches!(first[2], LogEntry::AnnotatedAssistantItem { .. })); + let first_bytes = fs::read(&path).unwrap(); + drop(store); + + let reopened = WorkerSessionStore::new(root.path()).unwrap(); + assert_eq!(fs::read(&path).unwrap(), first_bytes); + let snapshot = crate::public_snapshot::project_current_session_snapshot( + &reopened.read_all(session_id, segment_id).unwrap(), + ); + assert_eq!(snapshot.entries.len(), 3); + assert!(snapshot.entries.iter().all(|entry| { + entry.provenance == protocol::SessionEntryProvenance::LegacyUnknown + && entry.entry_id.len() <= 64 + })); + } + + #[test] + fn schema_v3_rejects_legacy_records_and_new_writes_are_canonical() { + let root = tempfile::tempdir().unwrap(); + let session_id = new_session_id(); + let segment_id = new_segment_id(); + let store = WorkerSessionStore::new(root.path()).unwrap(); + store + .create_segment( + session_id, + segment_id, + &[LogEntry::SegmentStart { + ts: 1, + session_id, + system_prompt: None, + config: agen::llm_client::RequestConfig::default(), + history: Vec::new(), + forked_from: None, + compacted_from: None, + }], + ) + .unwrap(); + store + .append( + session_id, + segment_id, + &LogEntry::UserInput { + ts: 2, + segments: vec![protocol::Segment::Text { + content: "new".into(), + }], + extensions: Vec::new(), + }, + ) + .unwrap(); + let entries = store.read_all(session_id, segment_id).unwrap(); + assert!(matches!(entries[0], LogEntry::AnnotatedSegmentStart { .. })); + assert!(matches!(entries[1], LogEntry::AnnotatedUserInput { .. })); + drop(store); + + let path = root + .path() + .join(SEGMENTS_DIR) + .join(format!("{segment_id}.jsonl")); + let mut file = OpenOptions::new().append(true).open(path).unwrap(); + serde_json::to_writer( + &mut file, + &LogEntry::SystemItem { + ts: 3, + item: crate::SystemItem::LegacyIgnored { + slug: "legacy".into(), + }, + }, + ) + .unwrap(); + file.write_all(b"\n").unwrap(); + let error = match WorkerSessionStore::new(root.path()) { + Ok(_) => panic!("schema v3 must reject a legacy history record"), + Err(error) => error, + }; + assert!(matches!(error, StoreError::Corrupt { .. })); + } + #[test] fn reopen_preserves_session_and_segment_ids() { let root = tempfile::tempdir().unwrap(); diff --git a/crates/tui/src/app.rs b/crates/tui/src/app.rs index fd8b7bc9..ba158e8b 100644 --- a/crates/tui/src/app.rs +++ b/crates/tui/src/app.rs @@ -1098,10 +1098,9 @@ impl App { self.blocks.push(Block::UserMessage { segments }); self.assistant_streaming = false; } - Event::SegmentRotated { entry } => { + Event::SegmentRotated { session } => { let retained_run_errors = self.run_error_messages.clone(); - self.reset_for_rotation(); - self.apply_log_entry_raw(&entry); + self.restore_session(&session, self.greeting.clone()); for message in retained_run_errors { self.blocks.push(Block::Alert { level: AlertLevel::Error, @@ -1408,14 +1407,14 @@ impl App { self.latest_memory_worker_event = Some(event.message); } Event::Snapshot { - entries, + session, greeting, status, in_flight, internal_workers, } => { self.rewind_refresh_fence = false; - self.restore_snapshot(&entries, greeting, in_flight); + self.restore_snapshot(&session, greeting, in_flight); self.replace_internal_worker_snapshots(internal_workers); self.set_worker_status(status); } @@ -1455,11 +1454,11 @@ impl App { } } Event::RewindApplied { - entries, + session, input, summary, } => { - self.restore_rewind_snapshot(&entries); + self.restore_rewind_snapshot(&session); self.rewind_refresh_fence = true; let restored_composer = if self.input.is_empty() { self.input.replace_with_segments(&input); @@ -2173,7 +2172,7 @@ impl App { ) -> InternalWorkerView { let mut app = App::new(snapshot.worker.name.clone()); app.mode = mode; - app.restore_entries(&snapshot.entries, None); + app.restore_session(&snapshot.session, None); app.apply_in_flight_snapshot(snapshot.in_flight); app.set_worker_status(snapshot.status); if let Some(error) = snapshot.error { @@ -2254,14 +2253,14 @@ impl App { fn restore_snapshot( &mut self, - entries: &[serde_json::Value], + session: &protocol::SessionSnapshot, greeting: protocol::Greeting, in_flight: InFlightSnapshot, ) { self.greeting = Some(greeting.clone()); self.context_window = greeting.context_window; self.session_context_tokens = greeting.context_tokens; - self.restore_entries(entries, Some(greeting)); + self.restore_session(session, Some(greeting)); self.apply_in_flight_snapshot(in_flight); } @@ -2270,7 +2269,7 @@ impl App { /// session tail; always clear/replay from it even if this TUI instance has /// somehow lost connect-time greeting metadata. Skipping the restore in /// that case would leave old post-target output visible after success. - fn restore_rewind_snapshot(&mut self, entries: &[serde_json::Value]) { + fn restore_rewind_snapshot(&mut self, session: &protocol::SessionSnapshot) { let greeting = self.greeting.clone().or_else(|| { self.blocks.iter().find_map(|b| match b { Block::Greeting(g) => Some(g.clone()), @@ -2283,7 +2282,7 @@ impl App { self.session_context_tokens = greeting.context_tokens; } let missing_greeting = greeting.is_none(); - self.restore_entries(entries, greeting); + self.restore_session(session, greeting); if missing_greeting { self.blocks.push(Block::Alert { level: AlertLevel::Warn, @@ -2293,9 +2292,9 @@ impl App { } } - fn restore_entries( + fn restore_session( &mut self, - entries: &[serde_json::Value], + session: &protocol::SessionSnapshot, greeting: Option, ) { self.run_error_messages.clear(); @@ -2309,137 +2308,90 @@ impl App { } self.assistant_streaming = false; - for entry in entries { - self.apply_log_entry_raw(entry); + for entry in &session.entries { + use protocol::{SessionContentPart, SessionMessageRole, SessionSnapshotEntryData}; + match &entry.data { + SessionSnapshotEntryData::UserInput { segments } => { + self.turn_index += 1; + self.blocks.push(Block::TurnHeader { + turn: self.turn_index, + }); + if !segments.is_empty() { + self.blocks.push(Block::UserMessage { + segments: segments.clone(), + }); + } + } + SessionSnapshotEntryData::Message { role, content } => { + let role = match role { + SessionMessageRole::User => agen::Role::User, + SessionMessageRole::Assistant => agen::Role::Assistant, + }; + let item = agen::Item::Message { + id: None, + role, + content: content + .iter() + .map(|part| match part { + SessionContentPart::Text { text } => { + agen::ContentPart::Text { text: text.clone() } + } + SessionContentPart::Refusal { refusal } => { + agen::ContentPart::Refusal { + refusal: refusal.clone(), + } + } + }) + .collect(), + status: None, + }; + let value = serde_json::to_value(item).expect("Item is Serialize"); + self.push_history_item(&value); + } + SessionSnapshotEntryData::ToolCall { + call_id, + name, + arguments, + } => { + let item = + agen::Item::tool_call(call_id.clone(), name.clone(), arguments.clone()); + let value = serde_json::to_value(item).expect("Item is Serialize"); + self.push_history_item(&value); + } + SessionSnapshotEntryData::ToolResult { + call_id, + summary, + content, + is_error, + .. + } => { + let item = agen::Item::tool_result_item( + call_id.clone(), + summary.clone(), + content.clone(), + *is_error, + ); + let value = serde_json::to_value(item).expect("Item is Serialize"); + self.push_history_item(&value); + } + SessionSnapshotEntryData::SystemItem { data, .. } => { + if let Some(data) = data { + self.apply_system_item(data); + } + } + SessionSnapshotEntryData::RunError { message } => { + self.push_run_error(message.clone()); + } + } } - self.mark_orphan_tool_calls_incomplete_pass(); } - /// Drop the derived view in preparation for replaying a new - /// `SegmentStart` (compaction / fork). Greeting is preserved - /// because the Worker identity hasn't changed. - fn reset_for_rotation(&mut self) { - let greeting = self.blocks.iter().find_map(|b| match b { - Block::Greeting(g) => Some(g.clone()), - _ => None, - }); - self.turn_index = 0; - self.blocks.clear(); - self.cache = FileCache::new(); - self.task_store = TaskStore::new(); - self.task_pane_scroll = 0; - if let Some(g) = greeting { - self.greeting = Some(g.clone()); - self.blocks.push(Block::Greeting(g)); - } - } - - /// Walk a single `LogEntry` JSON value and translate it into blocks - /// the live event path would have produced. Shared between - /// `restore_snapshot` (replay path) and `apply_log_entry` (live - /// path). - fn apply_log_entry_raw(&mut self, value: &serde_json::Value) { - let Ok(entry) = serde_json::from_value::(value.clone()) else { - return; - }; - match entry { - session_store::LogEntry::SegmentStart { history, .. } => { - for logged in history { - let item: agen::Item = logged.into(); - let item_value = serde_json::to_value(&item).expect("Item is Serialize"); - self.push_history_item(&item_value); - } - } - session_store::LogEntry::UserInput { segments, .. } => { - self.turn_index += 1; - self.blocks.push(Block::TurnHeader { - turn: self.turn_index, - }); - if !segments.is_empty() { - self.blocks.push(Block::UserMessage { segments }); - } - } - session_store::LogEntry::AssistantItem { item, .. } - | session_store::LogEntry::ToolResult { item, .. } => { - let it: agen::Item = item.into(); - let item_value = serde_json::to_value(&it).expect("Item is Serialize"); - self.push_history_item(&item_value); - } - session_store::LogEntry::SystemItem { item, .. } => { - let value = serde_json::to_value(&item).expect("SystemItem is Serialize"); - self.apply_system_item(&value); - } - session_store::LogEntry::Extension { - domain, payload, .. - } if domain == "yoi.compaction" => { - self.apply_compaction_extension(&payload); - } - session_store::LogEntry::RunErrored { message, .. } => { - self.push_run_error(message); - } - // Non-history-bearing variants don't affect the block view. - _ => {} - } - } - /// Dispatch one `SystemItem` JSON value into the appropriate block. /// /// Kind-based routing replaces the old free-text `[Notification]` / /// `[File: …]` parsing path: each kind maps directly to a typed /// block (`Block::Notify`, `Block::WorkerEvent`, …). - fn apply_compaction_extension(&mut self, payload: &serde_json::Value) { - if payload.get("kind").and_then(|value| value.as_str()) != Some("compaction_block") { - return; - } - match payload.get("state").and_then(|value| value.as_str()) { - Some("running") => { - if self.last_streaming_compact_mut().is_none() { - self.blocks.push(Block::Compact(CompactEvent::Streaming { - started_at: Instant::now(), - })); - } - } - Some("done") => { - let new_segment_id = payload - .get("new_segment_id") - .and_then(|value| value.as_str()) - .and_then(|value| value.parse::().ok()) - .unwrap_or_else(uuid::Uuid::nil); - if let Some(evt) = self.last_streaming_compact_mut() { - *evt = CompactEvent::Done { - new_segment_id, - elapsed_secs: None, - }; - } else { - self.blocks.push(Block::Compact(CompactEvent::Done { - new_segment_id, - elapsed_secs: None, - })); - } - } - Some("failed") => { - let error = payload - .get("error") - .and_then(|value| value.as_str()) - .unwrap_or("compact failed") - .to_string(); - if let Some(evt) = self.last_streaming_compact_mut() { - *evt = CompactEvent::Failed { - error, - elapsed_secs: None, - }; - } else { - self.blocks.push(Block::Compact(CompactEvent::Failed { - error, - elapsed_secs: None, - })); - } - } - _ => {} - } - } - fn apply_system_item(&mut self, value: &serde_json::Value) { let Ok(item) = serde_json::from_value::(value.clone()) else { // Unknown / forward-compat shape: fall back to rendering the @@ -2542,6 +2494,15 @@ fn fmt_millis(ms: u64) -> String { } } +#[cfg(test)] +fn public_session(values: Vec) -> protocol::SessionSnapshot { + let entries = values + .into_iter() + .map(|value| serde_json::from_value(value).expect("LogEntry deserializes")) + .collect::>(); + session_store::public_snapshot::project_current_session_snapshot(&entries) +} + fn message_text(item: &serde_json::Value) -> String { item["content"] .as_array() @@ -2685,7 +2646,7 @@ mod rewind_refresh_tests { }); app.handle_worker_event(Event::RewindApplied { - entries: vec![], + session: protocol::SessionSnapshot { entries: vec![] }, input: vec![Segment::text("selected rewind input")], summary: summary(3), }); @@ -2704,7 +2665,7 @@ mod rewind_refresh_tests { }); app.handle_worker_event(Event::RewindApplied { - entries: vec![], + session: protocol::SessionSnapshot { entries: vec![] }, input: vec![Segment::text("rewound input")], summary: summary(1), }); @@ -2747,7 +2708,7 @@ mod rewind_refresh_tests { }); app.handle_worker_event(Event::RewindApplied { - entries: vec![], + session: protocol::SessionSnapshot { entries: vec![] }, input: vec![Segment::text("rewound input")], summary: summary(2), }); @@ -3289,7 +3250,9 @@ mod completion_flow_tests { }; app.handle_worker_event(Event::SegmentRotated { - entry: serde_json::to_value(start).expect("LogEntry is Serialize"), + session: public_session(vec![ + serde_json::to_value(start).expect("LogEntry is Serialize"), + ]), }); app.handle_worker_event(Event::UserMessage { segments: vec![Segment::text("first persisted message")], @@ -3533,7 +3496,7 @@ mod completion_flow_tests { } #[test] - fn snapshot_renders_system_message_block_from_session_start() { + fn snapshot_excludes_system_prompt_history_from_public_blocks() { let mut app = App::new("test".into()); let session_start = session_store::LogEntry::SegmentStart { ts: 1, @@ -3549,7 +3512,7 @@ mod completion_flow_tests { let session_start_value = serde_json::to_value(&session_start).unwrap(); app.handle_worker_event(Event::Snapshot { greeting: test_greeting(), - entries: vec![session_start_value], + session: public_session(vec![session_start_value]), status: WorkerStatus::Running, in_flight: Default::default(), internal_workers: Vec::new(), @@ -3557,10 +3520,8 @@ mod completion_flow_tests { assert!(matches!(app.worker_status, WorkerStatus::Running)); assert!(app.running); - assert!(matches!( - app.blocks.get(1), - Some(Block::SystemMessage { text }) if text == "[File: src/main.rs]\nfn main() {}" - )); + assert_eq!(app.blocks.len(), 1); + assert!(matches!(app.blocks.first(), Some(Block::Greeting(_)))); } #[test] @@ -3595,7 +3556,7 @@ mod completion_flow_tests { }; app.handle_worker_event(Event::Snapshot { greeting: test_greeting(), - entries: vec![serde_json::to_value(run_errored).unwrap()], + session: public_session(vec![serde_json::to_value(run_errored).unwrap()]), status: WorkerStatus::Idle, in_flight: Default::default(), internal_workers: Vec::new(), @@ -3633,7 +3594,7 @@ mod completion_flow_tests { compacted_from: None, }; app.handle_worker_event(Event::SegmentRotated { - entry: serde_json::to_value(segment_start).unwrap(), + session: public_session(vec![serde_json::to_value(segment_start).unwrap()]), }); let errors = app @@ -3656,7 +3617,9 @@ mod completion_flow_tests { let mut app = App::new("test".into()); app.handle_worker_event(Event::Snapshot { greeting: test_greeting(), - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, status: WorkerStatus::Running, in_flight: InFlightSnapshot { blocks: vec![ @@ -3762,7 +3725,9 @@ mod completion_flow_tests { }, revision, status: WorkerStatus::Idle, - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, in_flight: protocol::InFlightSnapshot::default(), error: None, internal_workers: Vec::new(), @@ -3977,7 +3942,9 @@ mod completion_flow_tests { assert_eq!(app.selected_worker_view().worker_name, "parent"); app.handle_worker_event(Event::Snapshot { greeting: test_greeting(), - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, status: WorkerStatus::Idle, in_flight: Default::default(), internal_workers: Vec::new(), @@ -4026,7 +3993,9 @@ mod completion_flow_tests { }); app.handle_worker_event(Event::Snapshot { greeting: test_greeting(), - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, status: WorkerStatus::Idle, in_flight: Default::default(), internal_workers: vec![InternalWorkerSnapshot { @@ -4037,7 +4006,9 @@ mod completion_flow_tests { kind: protocol::InternalWorkerKind::SubWorker, }, revision: 4, - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, status: WorkerStatus::Running, error: None, in_flight: Default::default(), @@ -4193,7 +4164,9 @@ mod completion_flow_tests { greeting.context_tokens = 45_000; app.handle_worker_event(Event::Snapshot { - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, greeting, status: WorkerStatus::Idle, in_flight: Default::default(), @@ -4396,7 +4369,7 @@ mod completion_flow_tests { ]; app.handle_worker_event(Event::Snapshot { greeting: test_greeting(), - entries: assistant_item_entries, + session: public_session(assistant_item_entries), status: WorkerStatus::Running, in_flight: Default::default(), internal_workers: Vec::new(), diff --git a/crates/tui/src/console/mod.rs b/crates/tui/src/console/mod.rs index 0a6262ee..92a3a193 100644 --- a/crates/tui/src/console/mod.rs +++ b/crates/tui/src/console/mod.rs @@ -547,7 +547,9 @@ async fn run_e2e_rewind_fixture( let mut app = App::new_with_persistent_input_history(worker_name.clone(), &workspace_root); app.connected = true; app.handle_worker_event(Event::Snapshot { - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, status: WorkerStatus::Idle, greeting: Greeting { worker_name: worker_name.clone(), @@ -673,7 +675,9 @@ async fn run_e2e_rewind_fixture( if let Some(submitted_at) = pending_apply { if submitted_at.elapsed() >= apply_delay { app.handle_worker_event(Event::RewindApplied { - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, input: vec![Segment::text("rewind-live-refresh")], summary: RewindSummary { truncated_to_entries: 1, @@ -2023,13 +2027,13 @@ mod tests { let mut app = App::new("agent".to_string()); app.handle_worker_event(Event::Snapshot { greeting: test_greeting(), - entries: vec![], + session: protocol::SessionSnapshot { entries: vec![] }, status: WorkerStatus::Idle, in_flight: Default::default(), internal_workers: Vec::new(), }); app.handle_worker_event(Event::RewindApplied { - entries: vec![], + session: protocol::SessionSnapshot { entries: vec![] }, input: vec![Segment::Text { content: "retry this".into(), }], @@ -2050,7 +2054,7 @@ mod tests { let mut app = App::new("agent".to_string()); app.handle_worker_event(Event::Snapshot { greeting: test_greeting(), - entries: vec![], + session: protocol::SessionSnapshot { entries: vec![] }, status: WorkerStatus::Idle, in_flight: Default::default(), internal_workers: Vec::new(), @@ -2058,7 +2062,7 @@ mod tests { type_keys(&mut app, "draft"); app.handle_worker_event(Event::RewindApplied { - entries: vec![], + session: protocol::SessionSnapshot { entries: vec![] }, input: vec![Segment::Text { content: "retry this".into(), }], diff --git a/crates/tui/src/dashboard/tests.rs b/crates/tui/src/dashboard/tests.rs index 32be9393..a86a3040 100644 --- a/crates/tui/src/dashboard/tests.rs +++ b/crates/tui/src/dashboard/tests.rs @@ -849,7 +849,9 @@ async fn ticket_queue_notification_sends_notify_when_socket_available() { let mut writer = JsonLineWriter::new(writer); writer .write(&Event::Snapshot { - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, greeting: protocol::Greeting { worker_name: "test-orchestrator".to_string(), cwd: temp.path().display().to_string(), @@ -891,7 +893,9 @@ async fn send_notify_only_can_deliver_weak_notification_without_auto_run() { let mut writer = JsonLineWriter::new(writer); writer .write(&Event::Snapshot { - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, greeting: protocol::Greeting { worker_name: "yoi".to_string(), cwd: temp.path().display().to_string(), diff --git a/crates/tui/src/worker_list.rs b/crates/tui/src/worker_list.rs index 49236489..fa2d834c 100644 --- a/crates/tui/src/worker_list.rs +++ b/crates/tui/src/worker_list.rs @@ -910,7 +910,7 @@ mod tests { timestamp_ms: 0, }), Event::Snapshot { - entries: vec![], + session: protocol::SessionSnapshot { entries: vec![] }, greeting: test_greeting(), status: WorkerStatus::Idle, in_flight: Default::default(), diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 03f3b5f4..a1d2a99d 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -1530,7 +1530,9 @@ impl Runtime { } } Ok(protocol::Event::Snapshot { - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, greeting: protocol::Greeting { worker_name: worker_ref.worker_id.to_string(), cwd: String::new(), @@ -3152,7 +3154,9 @@ mod tests { ), ); let snapshot = protocol::Event::Snapshot { - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, greeting: protocol::Greeting { worker_name: "parent".to_string(), cwd: "/tmp".to_string(), @@ -4581,7 +4585,16 @@ mod tests { backend.set_worker_snapshot( &detail.worker_ref, protocol::Event::Snapshot { - entries: vec![expected_entry.clone()], + session: protocol::SessionSnapshot { + entries: vec![protocol::SessionSnapshotEntry { + entry_id: "restored-log-entry".to_owned(), + provenance: protocol::SessionEntryProvenance::LegacyUnknown, + derived_from: Vec::new(), + data: protocol::SessionSnapshotEntryData::RunError { + message: expected_entry.to_string(), + }, + }], + }, greeting: protocol::Greeting { worker_name: "live-worker".to_string(), cwd: "/tmp/live".to_string(), @@ -4606,12 +4619,13 @@ mod tests { .unwrap(); match snapshot { protocol::Event::Snapshot { - entries, + session, greeting, status, .. } => { - assert_eq!(entries, vec![expected_entry]); + assert_eq!(session.entries.len(), 1); + assert_eq!(session.entries[0].entry_id, "restored-log-entry"); assert_eq!(greeting.worker_name, "live-worker"); assert_eq!(status, protocol::WorkerStatus::Running); } diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index ff943c60..4042ab68 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -84,10 +84,7 @@ impl WorkerHandle { (entries, entry_rx, in_flight) }; let event = Event::Snapshot { - entries: entries - .into_iter() - .map(|entry| serde_json::to_value(entry).expect("log entry serializes")) - .collect(), + session: session_store::public_snapshot::project_current_session_snapshot(&entries), greeting: self.shared_state.greeting.clone(), status: self.shared_state.get_status(), in_flight, @@ -1874,28 +1871,16 @@ where St: Store, { match worker.rewind_to(target, expected_head_entries) { - Ok(applied) => match applied - .entries - .into_iter() - .map(serde_json::to_value) - .collect::, _>>() - { - Ok(entries) => { - let _ = event_tx.send(Event::RewindApplied { - entries, - input: applied.input, - summary: applied.summary, - }); - true - } - Err(error) => { - let _ = event_tx.send(Event::Error { - code: ErrorCode::Internal, - message: format!("failed to encode rewind snapshot: {error}"), - }); - false - } - }, + Ok(applied) => { + let session = + session_store::public_snapshot::project_current_session_snapshot(&applied.entries); + let _ = event_tx.send(Event::RewindApplied { + session, + input: applied.input, + summary: applied.summary, + }); + true + } Err(err) => { let _ = event_tx.send(Event::Error { code: ErrorCode::InvalidRequest, @@ -2101,7 +2086,9 @@ mod tests { let mut writer = JsonLineWriter::new(w); writer .write(&Event::Snapshot { - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, greeting: protocol::Greeting { worker_name: "parent".into(), cwd: "/tmp".into(), diff --git a/crates/worker/src/discovery.rs b/crates/worker/src/discovery.rs index 6e31005b..03a499d8 100644 --- a/crates/worker/src/discovery.rs +++ b/crates/worker/src/discovery.rs @@ -1481,7 +1481,9 @@ mod tests { let mut writer = JsonLineWriter::new(stream); writer .write(&Event::Snapshot { - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, greeting: protocol::Greeting { worker_name: "target".into(), cwd: "/tmp".into(), @@ -1514,7 +1516,9 @@ mod tests { .unwrap(); writer .write(&Event::Snapshot { - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, greeting: protocol::Greeting { worker_name: "target".into(), cwd: "/tmp".into(), @@ -1603,7 +1607,9 @@ mod tests { let mut writer = JsonLineWriter::new(stream); writer .write(&Event::Snapshot { - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, greeting: protocol::Greeting { worker_name: "target".into(), cwd: "/tmp".into(), @@ -1627,7 +1633,9 @@ mod tests { let mut writer = JsonLineWriter::new(writer_half); writer .write(&Event::Snapshot { - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, greeting: protocol::Greeting { worker_name: "target".into(), cwd: "/tmp".into(), @@ -1729,7 +1737,9 @@ mod tests { .unwrap(); writer .write(&Event::Snapshot { - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, greeting: protocol::Greeting { worker_name: "alerted".into(), cwd: "/tmp".into(), @@ -1779,7 +1789,9 @@ mod tests { let mut writer = JsonLineWriter::new(stream); let _ = writer .write(&Event::Snapshot { - entries: Vec::new(), + session: protocol::SessionSnapshot { + entries: Vec::new(), + }, greeting: protocol::Greeting { worker_name: "child-live".into(), cwd: "/tmp".into(), diff --git a/crates/worker/src/feature/builtin/worker_observation.rs b/crates/worker/src/feature/builtin/worker_observation.rs index 68305b2c..ad59bb3f 100644 --- a/crates/worker/src/feature/builtin/worker_observation.rs +++ b/crates/worker/src/feature/builtin/worker_observation.rs @@ -6,7 +6,7 @@ use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use async_trait::async_trait; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; -use session_store::{LogEntry, collect_state}; +use session_store::LogEntry; use super::manage_worker::{WORKER_CONTROL_SERVICE_ID, WorkerControlService}; use crate::feature::{ @@ -61,7 +61,7 @@ pub struct WorkerObservationSubject { #[derive(Debug, Clone)] pub struct WorkerSessionCapture { pub segment_id: String, - pub entries: Vec>, + pub session: protocol::SessionSnapshot, } impl WorkerSessionCapture { @@ -69,17 +69,9 @@ impl WorkerSessionCapture { segment_id: impl Into, log_entries: &[LogEntry], ) -> Result { - let segment_id = segment_id.into(); - let state = collect_state(log_entries); - let parsed_segment_id = segment_id.parse().unwrap_or_default(); - let entries = crate::session_history::restore_history_entries( - state.session_id.unwrap_or_default(), - parsed_segment_id, - log_entries, - )?; Ok(Self { - segment_id, - entries, + segment_id: segment_id.into(), + session: session_store::public_snapshot::project_current_session_snapshot(log_entries), }) } } @@ -115,7 +107,7 @@ struct WorkspaceWorkerObservationListResponse { #[derive(Debug, Deserialize)] struct WorkspaceWorkerObservationCaptureResponse { segment_id: String, - entries: Vec, + session: protocol::SessionSnapshot, } pub struct WorkspaceClientWorkerObservationProvider { @@ -173,26 +165,9 @@ impl WorkerObservationProvider for WorkspaceClientWorkerObservationProvider { let body = workspace_response_body(response)?; let response = serde_json::from_str::(&body) .map_err(|error| WorkerObservationError::Unavailable(error.to_string()))?; - let entries = response - .entries - .into_iter() - .map(|entry| { - serde_json::from_value(entry) - .map_err(|error| WorkerObservationError::Unavailable(error.to_string())) - }) - .collect::, _>>()?; - let state = collect_state(&entries); - let segment_id = response.segment_id; - let parsed_segment_id = segment_id.parse().unwrap_or_default(); - let typed_entries = crate::session_history::restore_history_entries( - state.session_id.unwrap_or_default(), - parsed_segment_id, - &entries, - ) - .map_err(WorkerObservationError::Unavailable)?; Ok(WorkerSessionCapture { - segment_id, - entries: typed_entries, + segment_id: response.segment_id, + session: response.session, }) } } @@ -420,16 +395,9 @@ impl WorkerObservationProvider for SpawnedSubWorkerObservationProvider { .get_internal(name) .ok_or(WorkerObservationError::NotFound)?; let entries = record.session.entries(); - let state = collect_state(&entries); - let typed_entries = crate::session_history::restore_history_entries( - state.session_id.unwrap_or_default(), - Default::default(), - &entries, - ) - .map_err(WorkerObservationError::Unavailable)?; Ok(WorkerSessionCapture { segment_id: format!("subworker:{name}"), - entries: typed_entries, + session: session_store::public_snapshot::project_current_session_snapshot(&entries), }) } } @@ -699,9 +667,9 @@ async fn latest_view( .capture_worker_session(subject) .await .map_err(tool_error)?; - Ok(SessionCapture::from_history_entries( + Ok(SessionCapture::from_session_snapshot( capture.segment_id, - capture.entries, + capture.session, )) } @@ -799,16 +767,42 @@ mod tests { .clone() .into_iter() .enumerate() - .map(|(index, item)| { - let mut metadata = crate::SessionHistoryMetadata::legacy_unknown(); - metadata.entry_id = - session_store::LoggedSessionHistoryEntryId(format!("fake-{index:08}")); - agen::HistoryEntry::new(item, metadata) + .filter_map(|(index, item)| { + let data = match item { + Item::Message { role, content, .. } => { + let role = match role { + Role::User => protocol::SessionMessageRole::User, + Role::Assistant => protocol::SessionMessageRole::Assistant, + Role::System => return None, + }; + protocol::SessionSnapshotEntryData::Message { + role, + content: content + .into_iter() + .map(|part| match part { + agen::ContentPart::Text { text } => { + protocol::SessionContentPart::Text { text } + } + agen::ContentPart::Refusal { refusal } => { + protocol::SessionContentPart::Refusal { refusal } + } + }) + .collect(), + } + } + _ => return None, + }; + Some(protocol::SessionSnapshotEntry { + entry_id: format!("fake-{index:08}"), + provenance: protocol::SessionEntryProvenance::LegacyUnknown, + derived_from: Vec::new(), + data, + }) }) .collect(); Ok(WorkerSessionCapture { segment_id: "segment".to_string(), - entries, + session: protocol::SessionSnapshot { entries }, }) } } diff --git a/crates/worker/src/internal_worker.rs b/crates/worker/src/internal_worker.rs index a8d21e49..e033369e 100644 --- a/crates/worker/src/internal_worker.rs +++ b/crates/worker/src/internal_worker.rs @@ -335,7 +335,7 @@ enum InternalWorkerSessionCommand { /// task; protocol access is consumed only by the owning parent registry. #[derive(Debug, Clone)] pub(crate) struct InternalWorkerSessionSnapshot { - pub entries: Vec, + pub session: protocol::SessionSnapshot, pub status: WorkerStatus, pub error: Option, pub in_flight: InFlightSnapshot, @@ -402,7 +402,7 @@ impl InternalWorkerSessionHandle { (entries, snapshot_from_guard(&guard)) }; InternalWorkerSessionSnapshot { - entries, + session: session_store::public_snapshot::project_current_session_snapshot(&entries), status: match self.status() { InternalWorkerSessionStatus::Running => WorkerStatus::Running, InternalWorkerSessionStatus::Paused => WorkerStatus::Paused, diff --git a/crates/worker/src/ipc/protocol_session.rs b/crates/worker/src/ipc/protocol_session.rs index 9b749126..5e16069b 100644 --- a/crates/worker/src/ipc/protocol_session.rs +++ b/crates/worker/src/ipc/protocol_session.rs @@ -30,8 +30,9 @@ pub fn subscribe_worker_protocol_session(handle: &WorkerHandle) -> WorkerProtoco pub fn live_log_entry_event(entry: LogEntry) -> Option { match entry { entry @ (LogEntry::SegmentStart { .. } | LogEntry::AnnotatedSegmentStart { .. }) => { - let value = serde_json::to_value(&entry).expect("LogEntry is Serialize"); - Some(Event::SegmentRotated { entry: value }) + let session = + session_store::public_snapshot::project_current_session_snapshot(&[entry]); + Some(Event::SegmentRotated { session }) } LogEntry::UserInput { segments, .. } | LogEntry::AnnotatedUserInput { segments, .. } => { Some(Event::UserMessage { segments }) diff --git a/crates/worker/src/session_capture.rs b/crates/worker/src/session_capture.rs index 1886e4e8..c0be48d0 100644 --- a/crates/worker/src/session_capture.rs +++ b/crates/worker/src/session_capture.rs @@ -8,6 +8,10 @@ use std::sync::Arc; use crate::session_history::{SessionHistoryMetadata, WorkerHistoryProvenance}; use agen::{HistoryEntry, Item, Role}; +use protocol::{ + SessionContentPart, SessionEntryProvenance, SessionMessageRole, SessionSnapshot, + SessionSnapshotEntryData, +}; use serde::{Deserialize, Serialize}; const DEFAULT_SEARCH_LIMIT: usize = 20; @@ -225,6 +229,66 @@ pub(crate) struct SessionCapture { } impl SessionCapture { + pub(crate) fn from_session_snapshot( + segment_id: impl Into, + snapshot: SessionSnapshot, + ) -> Self { + let entries = snapshot + .entries + .into_iter() + .filter_map(|entry| { + let item = match entry.data { + SessionSnapshotEntryData::UserInput { segments } => { + Item::user_message(protocol::Segment::flatten_to_text(&segments)) + } + SessionSnapshotEntryData::Message { role, content } => { + let role = match role { + SessionMessageRole::User => Role::User, + SessionMessageRole::Assistant => Role::Assistant, + }; + Item::Message { + id: None, + role, + content: content + .into_iter() + .map(|part| match part { + SessionContentPart::Text { text } => { + agen::ContentPart::Text { text } + } + SessionContentPart::Refusal { refusal } => { + agen::ContentPart::Refusal { refusal } + } + }) + .collect(), + status: None, + } + } + SessionSnapshotEntryData::ToolCall { + call_id, + name, + arguments, + } => Item::tool_call(call_id, name, arguments), + SessionSnapshotEntryData::ToolResult { + call_id, + summary, + content, + is_error, + attachments: _, + } => Item::tool_result_item(call_id, summary, content, is_error), + // Observation deliberately excludes system items and + // controller errors from model-visible session evidence. + SessionSnapshotEntryData::SystemItem { .. } + | SessionSnapshotEntryData::RunError { .. } => return None, + }; + Some(HistoryEntry::new( + item, + public_snapshot_metadata(entry.entry_id, entry.provenance), + )) + }) + .collect(); + Self::from_history_entries(segment_id, entries) + } + pub(crate) fn new(segment_id: impl Into, items: Vec) -> Self { let entries = items .into_iter() @@ -543,6 +607,46 @@ impl SessionCapture { } } +fn public_snapshot_metadata( + entry_id: String, + provenance: SessionEntryProvenance, +) -> SessionHistoryMetadata { + let worker = session_store::LoggedWorkerSubject { + workspace_id: None, + runtime_id: None, + 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 { + worker: worker.clone(), + }, + SessionEntryProvenance::ToolOutput => WorkerHistoryProvenance::ToolOutput { worker }, + SessionEntryProvenance::DerivedSummary => WorkerHistoryProvenance::DerivedSummary, + SessionEntryProvenance::LegacyUnknown => WorkerHistoryProvenance::LegacyUnknown, + }; + SessionHistoryMetadata { + entry_id: session_store::LoggedSessionHistoryEntryId(entry_id), + origin, + derivation: None, + } +} + fn message_reference_kind( origin: &WorkerHistoryProvenance, provider_role: &Role, diff --git a/crates/worker/src/spawn/comm_tools.rs b/crates/worker/src/spawn/comm_tools.rs index 03b1b126..9d30d89b 100644 --- a/crates/worker/src/spawn/comm_tools.rs +++ b/crates/worker/src/spawn/comm_tools.rs @@ -278,7 +278,20 @@ mod tests { fn snapshot(entries: Vec) -> Event { Event::Snapshot { - entries, + session: protocol::SessionSnapshot { + entries: entries + .into_iter() + .enumerate() + .map(|(index, value)| protocol::SessionSnapshotEntry { + entry_id: format!("test-{index}"), + provenance: protocol::SessionEntryProvenance::LegacyUnknown, + derived_from: Vec::new(), + data: protocol::SessionSnapshotEntryData::RunError { + message: value.to_string(), + }, + }) + .collect(), + }, greeting: Greeting { worker_name: "server".into(), cwd: "/tmp".into(), diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index 86ea8bae..8e0ea0f1 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -806,11 +806,7 @@ fn internal_worker_snapshot( InternalWorkerSnapshot { worker, revision, - entries: snapshot - .entries - .into_iter() - .filter_map(|entry| serde_json::to_value(entry).ok()) - .collect(), + session: snapshot.session, status: snapshot.status, error: snapshot.error, in_flight: snapshot.in_flight, @@ -1045,7 +1041,7 @@ mod tests { let snapshots = registry.internal_worker_snapshots(); assert_eq!(snapshots.len(), 1); assert_eq!(snapshots[0].revision, 2); - assert_eq!(snapshots[0].entries.len(), 1); + assert_eq!(snapshots[0].session.entries.len(), 1); record.session.emit_test_text_delta("partial"); let streamed = tokio::time::timeout(Duration::from_secs(1), parent_rx.recv()) diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index 21a9371a..8e4843d7 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -960,9 +960,7 @@ mod tests { use crate::WorkspaceId; use agen::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent}; - use agen::llm_client::types::ContentPart; use agen::llm_client::{ClientError, LlmClient, Request}; - use agen::{Item, Role}; use async_trait::async_trait; use futures::Stream; use manifest::{AuthRef, ModelManifest, SchemeKind, WorkerManifest}; @@ -1252,9 +1250,11 @@ extract_threshold = 4000 ) .await .unwrap(); - assert!(first_capture.entries.iter().map(|entry| &entry.item).any(|item| { - matches!(item, Item::Message { role: Role::Assistant, content, .. } if content.iter().any(|part| matches!(part, ContentPart::Text { text } if text.contains("reviewed")))) - })); + assert!( + serde_json::to_string(&first_capture.session) + .unwrap() + .contains("reviewed") + ); let send = (crate::spawn::comm_tools::sub_worker_send_tool(registry.clone()))().1; send.execute( @@ -1274,7 +1274,7 @@ extract_threshold = 4000 ) .await .unwrap(); - assert!(latest_capture.entries.len() > first_capture.entries.len()); + assert!(latest_capture.session.entries.len() > first_capture.session.entries.len()); fail_requests.store(true, Ordering::SeqCst); send.execute( diff --git a/crates/worker/tests/controller_test.rs b/crates/worker/tests/controller_test.rs index e4ef5b31..078fdc97 100644 --- a/crates/worker/tests/controller_test.rs +++ b/crates/worker/tests/controller_test.rs @@ -839,26 +839,26 @@ async fn snapshot_includes_user_input_for_in_flight_turn() { loop { let event = reader.next::().await.unwrap().unwrap(); match event { - Event::Snapshot { entries, .. } => { - // Walk the entries, find a `LogEntry::UserInput` and - // confirm its segments flatten to our submitted text. - let mut found = false; - for value in &entries { - let entry: session_store::LogEntry = - serde_json::from_value(value.clone()).expect("LogEntry deserialise"); - if let session_store::LogEntry::UserInput { segments, .. } - | session_store::LogEntry::AnnotatedUserInput { segments, .. } = entry - { - let text = protocol::Segment::flatten_to_text(&segments); - if text == "hello in-flight" { - found = true; - break; - } + Event::Snapshot { session, .. } => { + let found = session.entries.iter().any(|entry| match &entry.data { + protocol::SessionSnapshotEntryData::UserInput { segments } => { + protocol::Segment::flatten_to_text(segments) == "hello in-flight" } - } + protocol::SessionSnapshotEntryData::Message { + role: protocol::SessionMessageRole::User, + content, + } => content.iter().any(|part| { + matches!( + part, + protocol::SessionContentPart::Text { text } + if text == "hello in-flight" + ) + }), + _ => false, + }); assert!( found, - "snapshot must carry the in-flight UserInput entry: {entries:?}" + "snapshot must carry the in-flight UserInput entry: {session:?}" ); return; } @@ -2410,17 +2410,12 @@ async fn snapshot_contains_user_input(handle: &WorkerHandle, needle: &str) -> bo loop { let event = reader.next::().await.unwrap().unwrap(); match event { - Event::Snapshot { entries, .. } => { - return entries.into_iter().any(|value| { - let entry: session_store::LogEntry = - serde_json::from_value(value).expect("LogEntry deserialise"); - match entry { - session_store::LogEntry::UserInput { segments, .. } - | session_store::LogEntry::AnnotatedUserInput { segments, .. } => { - protocol::Segment::flatten_to_text(&segments).contains(needle) - } - _ => false, + Event::Snapshot { session, .. } => { + return session.entries.into_iter().any(|entry| match entry.data { + protocol::SessionSnapshotEntryData::UserInput { segments } => { + protocol::Segment::flatten_to_text(&segments).contains(needle) } + _ => false, }); } Event::Alert(_) => continue, diff --git a/crates/worker/tests/system_prompt_template_test.rs b/crates/worker/tests/system_prompt_template_test.rs index 53ca4141..9293917d 100644 --- a/crates/worker/tests/system_prompt_template_test.rs +++ b/crates/worker/tests/system_prompt_template_test.rs @@ -203,12 +203,12 @@ async fn session_start_state_captures_rendered_prompt() { .unwrap(); let first = entries.first().expect("at least one entry"); match first { - LogEntry::SegmentStart { system_prompt, .. } => { + LogEntry::AnnotatedSegmentStart { system_prompt, .. } => { let sp = system_prompt.as_deref().expect("system prompt set"); assert!(sp.starts_with("hello")); assert!(sp.contains(&pwd.display().to_string())); } - other => panic!("expected SegmentStart as first entry, got {other:?}"), + other => panic!("expected AnnotatedSegmentStart as first entry, got {other:?}"), } } diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index cc374820..da87769a 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -8513,7 +8513,7 @@ async fn scoped_capture_worker_observation_session( message: "worker protocol closed before the session snapshot".to_string(), }) })?; - let protocol::Event::Snapshot { entries, .. } = event else { + let protocol::Event::Snapshot { session, .. } = event else { return Err(ApiError::from(Error::RuntimeOperationFailed { runtime_id: target.runtime_id.clone(), code: "worker_observation_missing_snapshot".to_string(), @@ -8522,7 +8522,7 @@ async fn scoped_capture_worker_observation_session( }; Ok(Json(serde_json::json!({ "segment_id": format!("runtime:{}:worker:{}", target.runtime_id, target.worker_id), - "entries": entries, + "session": session, }))) } @@ -16785,7 +16785,7 @@ mod tests { ) .await .unwrap(); - assert!(capture["entries"].is_array()); + assert!(capture["session"]["entries"].is_array()); let revoked = api .store diff --git a/web/workspace/src/lib/generated/protocol.ts b/web/workspace/src/lib/generated/protocol.ts index 6a357acb..39ab33c2 100644 --- a/web/workspace/src/lib/generated/protocol.ts +++ b/web/workspace/src/lib/generated/protocol.ts @@ -73,11 +73,35 @@ export type InFlightBlock = { "kind": "text", text: string, finished?: boolean, export type InFlightSnapshot = { blocks?: Array, commands?: Array, }; +export type SessionEntryProvenance = "human_input" | "worker_input" | "flow_instruction" | "backend_instruction" | "model_output" | "tool_output" | "derived_summary" | "legacy_unknown"; + +export type SessionMessageRole = "user" | "assistant"; + +export type SessionContentPart = { "kind": "text", text: string, } | { "kind": "refusal", refusal: string, }; + +export type SessionToolAttachment = { media_type: string, +/** + * Base64-encoded durable attachment body. Public snapshots preserve the + * committed multimodal value instead of replacing it with placeholder text. + */ +data_base64: string, }; + +export type SessionSnapshotEntryData = { "kind": "user_input", segments: Array, } | { "kind": "message", role: SessionMessageRole, content: Array, } | { "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, } | { "kind": "system_item", item_kind: string, content: string, data?: unknown, } | { "kind": "run_error", message: string, }; + +export type SessionSnapshotEntry = { +/** + * Stable identity from durable history metadata, or a deterministic + * identity derived from the legacy segment and log position. + */ +entry_id: string, provenance: SessionEntryProvenance, derived_from?: Array, } & ({ "kind": "user_input", segments: Array, } | { "kind": "message", role: SessionMessageRole, content: Array, } | { "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, } | { "kind": "system_item", item_kind: string, content: string, data?: unknown, } | { "kind": "run_error", message: string, }); + +export type SessionSnapshot = { entries: Array, }; + export type InternalWorkerKind = "sub_worker" | { "service": { kind: string, } }; export type InternalWorkerRef = { session_id: string, name: string, parent_session_id?: string | null, kind: InternalWorkerKind, }; -export type InternalWorkerSnapshot = { worker: InternalWorkerRef, revision: number, entries: Array, status: WorkerStatus, error?: string | null, in_flight?: InFlightSnapshot, internal_workers?: Array, }; +export type InternalWorkerSnapshot = { worker: InternalWorkerRef, revision: number, session: SessionSnapshot, status: WorkerStatus, error?: string | null, in_flight?: InFlightSnapshot, internal_workers?: Array, }; export type Greeting = { worker_name: string, cwd: string, provider: string, model: string, scope_summary: string, tools: Array, /** @@ -193,7 +217,7 @@ summary: string, * Full tool output. Absent when the tool chose to return * summary-only, or when the result was pruned. */ -output?: string | null, disposition?: ToolResultDisposition | null, is_error: boolean, } } | { "event": "usage", "data": { input_tokens: number | null, output_tokens: number | null, cache_read_input_tokens?: number | null, } } | { "event": "run_end", "data": { result: RunResult, } } | { "event": "error", "data": { code: ErrorCode, message: string, } } | { "event": "snapshot", "data": { entries: Array, greeting: Greeting, status: WorkerStatus, +output?: string | null, disposition?: ToolResultDisposition | null, is_error: boolean, } } | { "event": "usage", "data": { input_tokens: number | null, output_tokens: number | null, cache_read_input_tokens?: number | null, } } | { "event": "run_end", "data": { result: RunResult, } } | { "event": "error", "data": { code: ErrorCode, message: string, } } | { "event": "snapshot", "data": { session: SessionSnapshot, greeting: Greeting, status: WorkerStatus, /** * Unfinished model output that has already streamed in the current * run but is not yet represented by committed snapshot entries. @@ -203,4 +227,4 @@ in_flight?: InFlightSnapshot, * Parent-owned Internal Worker sessions visible to this client. * Service-private Internal Workers are deliberately excluded. */ -internal_workers?: Array, } } | { "event": "internal_worker", "data": { worker: InternalWorkerRef, revision: number, event: Event, } } | { "event": "internal_worker_removed", "data": { worker: InternalWorkerRef, revision: number, } } | { "event": "segment_rotated", "data": { entry: unknown, } } | { "event": "status", "data": { status: WorkerStatus, } } | { "event": "command", "data": { event: CommandEvent, } } | { "event": "completions", "data": { kind: CompletionKind, entries: Array, } } | { "event": "rewind_targets", "data": { head_entries: number, targets: Array, } } | { "event": "rewind_applied", "data": { entries: Array, input: Array, summary: RewindSummary, } } | { "event": "workers_listed", "data": { workers: unknown, } } | { "event": "worker_restored", "data": { result: unknown, } } | { "event": "peer_registered", "data": { result: unknown, } } | { "event": "alert", "data": Alert } | { "event": "memory_worker", "data": MemoryWorkerEvent } | { "event": "compact_start", "data": { lifecycle: CompactionLifecycle, } } | { "event": "compact_done", "data": { lifecycle: CompactionLifecycle, } } | { "event": "compact_failed", "data": { lifecycle: CompactionLifecycle, } } | { "event": "shutdown" }; +internal_workers?: Array, } } | { "event": "internal_worker", "data": { worker: InternalWorkerRef, revision: number, event: Event, } } | { "event": "internal_worker_removed", "data": { worker: InternalWorkerRef, revision: number, } } | { "event": "segment_rotated", "data": { session: SessionSnapshot, } } | { "event": "status", "data": { status: WorkerStatus, } } | { "event": "command", "data": { event: CommandEvent, } } | { "event": "completions", "data": { kind: CompletionKind, entries: Array, } } | { "event": "rewind_targets", "data": { head_entries: number, targets: Array, } } | { "event": "rewind_applied", "data": { session: SessionSnapshot, input: Array, summary: RewindSummary, } } | { "event": "workers_listed", "data": { workers: unknown, } } | { "event": "worker_restored", "data": { result: unknown, } } | { "event": "peer_registered", "data": { result: unknown, } } | { "event": "alert", "data": Alert } | { "event": "memory_worker", "data": MemoryWorkerEvent } | { "event": "compact_start", "data": { lifecycle: CompactionLifecycle, } } | { "event": "compact_done", "data": { lifecycle: CompactionLifecycle, } } | { "event": "compact_failed", "data": { lifecycle: CompactionLifecycle, } } | { "event": "shutdown" }; diff --git a/web/workspace/src/lib/workspace/console/model.test.ts b/web/workspace/src/lib/workspace/console/model.test.ts index 39732416..418b3d58 100644 --- a/web/workspace/src/lib/workspace/console/model.test.ts +++ b/web/workspace/src/lib/workspace/console/model.test.ts @@ -43,11 +43,84 @@ function consoleLine(id: string, kind: ConsoleLine["kind"]): ConsoleLine { }; } +function canonicalSession(logEntries: unknown[]): Event extends { + event: "snapshot"; + data: infer D; +} ? D extends { session: infer S } ? S : never : never { + const entries: Record[] = []; + let sequence = 0; + const itemEntry = (item: Record) => { + const kind = item["kind"]; + if (kind === "reasoning" || item["role"] === "system") return; + entries.push({ + entry_id: `legacy-test-${sequence++}`, + provenance: "legacy_unknown", + ...item, + }); + }; + for (const raw of logEntries) { + if (!raw || typeof raw !== "object" || Array.isArray(raw)) continue; + const entry = raw as Record; + switch (entry["kind"]) { + case "segment_start": + case "annotated_segment_start": + entries.length = 0; + for (const history of Array.isArray(entry["history"]) ? entry["history"] : []) { + const value = history as Record; + itemEntry((value["item"] as Record | undefined) ?? value); + } + break; + case "user_input": + case "annotated_user_input": + entries.push({ + entry_id: `legacy-test-${sequence++}`, + provenance: "legacy_unknown", + kind: "user_input", + segments: entry["segments"] ?? [], + }); + break; + case "assistant_item": + case "tool_result": + case "annotated_assistant_item": + case "annotated_tool_result": { + const annotated = entry["entry"] as Record | undefined; + itemEntry((annotated?.["item"] as Record | undefined) ?? + (entry["item"] as Record)); + break; + } + case "system_item": + case "annotated_system_item": { + const annotated = entry["entry"] as Record | undefined; + const item = (annotated?.["item"] as Record | undefined) ?? + (entry["item"] as Record); + entries.push({ + entry_id: `legacy-test-${sequence++}`, + provenance: "legacy_unknown", + kind: "system_item", + item_kind: item?.["kind"] ?? "system_item", + content: item?.["body"] ?? item?.["message"] ?? "", + data: item, + }); + break; + } + case "run_errored": + entries.push({ + entry_id: `legacy-test-${sequence++}`, + provenance: "legacy_unknown", + kind: "run_error", + message: entry["message"] ?? "Worker run failed.", + }); + break; + } + } + return { entries } as never; +} + function snapshotEvent(cwd: string, entries: unknown[] = []): Event { return { event: "snapshot", data: { - entries, + session: canonicalSession(entries), greeting: { worker_name: "Worker", cwd, @@ -138,7 +211,7 @@ Deno.test("segment rotation retains a live error beside the real SegmentStart hi event: { event: "segment_rotated", data: { - entry: { + session: canonicalSession([{ kind: "segment_start", ts: 5, session_id: "session-1", @@ -149,7 +222,7 @@ Deno.test("segment rotation retains a live error beside the real SegmentStart hi role: "user", content: [{ kind: "text", text: "retained conversation" }], }], - }, + }]), }, } satisfies Event, }, @@ -852,7 +925,7 @@ Deno.test("compaction service activity stays nested in one lifecycle item", () = ); }); -Deno.test("snapshot normalizes orphaned running compaction to interrupted", () => { +Deno.test("snapshot excludes storage-only compaction extension records", () => { const projection = projectConsole([{ eventId: "snapshot", observedAtMs: 9_000, @@ -876,10 +949,7 @@ Deno.test("snapshot normalizes orphaned running compaction to interrupted", () = }]), }]); - assertEquals(projection.lines.length, 1); - assertEquals(projection.lines[0].compaction?.state, "interrupted"); - assertEquals(projection.lines[0].compaction?.endedAtMs, 9_000); - assertEquals(projection.lines[0].streaming, false); + assertEquals(projection.lines.length, 0); }); Deno.test("createConsoleProjector ignores stale compaction revisions", () => { @@ -1238,7 +1308,7 @@ Deno.test("projectConsole renders snapshot entries and in-flight output", () => event: { event: "snapshot", data: { - entries: [ + session: canonicalSession([ { kind: "segment_start", ts: 1, @@ -1300,7 +1370,7 @@ Deno.test("projectConsole renders snapshot entries and in-flight output", () => message: "Compacting…", }, }, - ], + ]), greeting: { worker_name: "Worker", cwd: "/repo", @@ -1332,7 +1402,6 @@ Deno.test("projectConsole renders snapshot entries and in-flight output", () => "user:new user:false", "assistant:assistant reply:false", "tool:Read(1 file)\n /tmp/a.md:false", - "status:Compacting…:true", "in_flight:partial:true", ], ); @@ -1344,7 +1413,7 @@ Deno.test("projectConsole restores system items from snapshot entries", () => { event: { event: "snapshot", data: { - entries: [{ + session: canonicalSession([{ kind: "system_item", ts: 1, item: { @@ -1352,7 +1421,7 @@ Deno.test("projectConsole restores system items from snapshot entries", () => { message: "Worker completed", body: "Child Worker coder-1 completed.", }, - }], + }]), greeting: { worker_name: "Worker", cwd: "/repo", @@ -1388,7 +1457,7 @@ Deno.test("projectConsole reseeds visible rows from segment rotation", () => { event: { event: "segment_rotated", data: { - entry: { + session: canonicalSession([{ kind: "segment_start", ts: 10, session_id: "00000000-0000-0000-0000-000000000001", @@ -1401,7 +1470,7 @@ Deno.test("projectConsole reseeds visible rows from segment rotation", () => { content: [{ kind: "text", text: "after rotation seed" }], }, ], - }, + }]), }, } satisfies Event, }, @@ -1773,7 +1842,7 @@ Deno.test("parent snapshot authoritatively replaces Internal Worker projections" kind: "sub_worker", }, revision: 4, - entries: [{ + session: canonicalSession([{ kind: "assistant_item", ts: 1, item: { @@ -1792,7 +1861,7 @@ Deno.test("parent snapshot authoritatively replaces Internal Worker projections" content: "content", is_error: false, }, - }], + }]), status: "idle", in_flight: { blocks: [{ @@ -1934,18 +2003,20 @@ Deno.test("snapshot restores TaskStore state from system history", () => { `[Session TaskStore snapshot]\n\n\`\`\`json\n{\n "tasks": [{"taskid": 3, "status": "pending", "subject": "Restored", "description": "From compaction"}]\n}\n\`\`\``; const event = snapshotEvent("/repo"); if (event.event !== "snapshot") throw new Error("snapshot fixture expected"); - event.data.entries = [{ - kind: "segment_start", - ts: 1, - session_id: "00000000-0000-0000-0000-000000000001", - system_prompt: null, - config: {}, - history: [{ - kind: "message", - role: "system", - content: [{ kind: "text", text: taskSnapshot }], + event.data.session = { + entries: [{ + entry_id: "task-reminder-1", + provenance: "backend_instruction", + kind: "system_item", + item_kind: "task_reminder", + content: taskSnapshot, + data: { + kind: "task_reminder", + body: taskSnapshot, + source: "automatic", + }, }], - }]; + }; const projection = projectConsole([{ eventId: "task-snapshot", event }]); assertEquals(projection.tasks, [{ diff --git a/web/workspace/src/lib/workspace/console/model.ts b/web/workspace/src/lib/workspace/console/model.ts index d9e56953..61f34c78 100644 --- a/web/workspace/src/lib/workspace/console/model.ts +++ b/web/workspace/src/lib/workspace/console/model.ts @@ -650,9 +650,9 @@ function projectInternalWorkerSnapshot( eventId: string, cwd: string | null, ): InternalWorkerProjection { - const console = snapshotProjectionFromEntries( + const console = snapshotProjectionFromSession( `${eventId}:internal:${snapshot.worker.session_id}:snapshot`, - snapshot.entries, + snapshot.session, cwd, ); console.status = snapshot.status; @@ -905,9 +905,9 @@ export function applyProtocolEvent( case "snapshot": { next.status = event.data.status; next.cwd = event.data.greeting.cwd; - const snapshot = snapshotProjectionFromEntries( + const snapshot = snapshotProjectionFromSession( envelope.eventId, - event.data.entries, + event.data.session, next.cwd, ); next.lines = snapshot.lines; @@ -1008,9 +1008,9 @@ export function applyProtocolEvent( break; case "segment_rotated": { const retainedErrors = next.lines.filter((line) => line.kind === "error"); - const segment = snapshotProjectionFromEntries( + const segment = snapshotProjectionFromSession( envelope.eventId, - [event.data.entry], + event.data.session, next.cwd, ); next.lines = [...segment.lines, ...retainedErrors]; @@ -1925,9 +1925,9 @@ function applyTaskSystemItem( if (typeof body === "string") applyTaskSnapshot(projection, body); } -function snapshotProjectionFromEntries( +function snapshotProjectionFromSession( eventId: string, - entries: unknown[], + snapshot: unknown, cwd: string | null, ): ConsoleProjection { const projection: ConsoleProjection = { @@ -1942,58 +1942,74 @@ function snapshotProjectionFromEntries( internalWorkers: [], removedInternalWorkers: {}, }; + const entries = isRecord(snapshot) ? arrayField(snapshot, "entries") : []; entries.forEach((entry, index) => - applyLogEntry(projection, `${eventId}-snapshot-${index}`, entry) + applySessionEntry(projection, `${eventId}-snapshot-${index}`, entry) ); return projection; } -function applyLogEntry( +function applySessionEntry( projection: ConsoleProjection, - eventId: string, - entry: unknown, + fallbackEventId: string, + value: unknown, ): void { - if (!isRecord(entry)) return; - switch (stringField(entry, "kind")) { - case "segment_start": - arrayField(entry, "history").forEach((item, index) => - applyLoggedItem(projection, `${eventId}-history-${index}`, item) - ); - break; + if (!isRecord(value)) return; + const eventId = stringField(value, "entry_id") ?? fallbackEventId; + switch (stringField(value, "kind")) { case "user_input": projection.lines.push( line( eventId, "user", "User", - segmentsToText(arrayField(entry, "segments") as Segment[]), + segmentsToText(arrayField(value, "segments") as Segment[]), ), ); break; - case "system_item": - projection.lines.push(systemItemLine(eventId, entry["item"])); - applyTaskSystemItem(projection, entry["item"]); + case "message": + applyLoggedItem(projection, eventId, { + kind: "message", + role: value["role"], + content: value["content"], + }); + break; + case "tool_call": + applyLoggedItem(projection, eventId, { + kind: "tool_call", + call_id: value["call_id"], + name: value["name"], + arguments: value["arguments"], + }); break; - case "assistant_item": case "tool_result": - applyLoggedItem(projection, eventId, entry["item"]); + applyLoggedItem(projection, eventId, { + kind: "tool_result", + call_id: value["call_id"], + summary: value["summary"], + content: value["content"], + is_error: value["is_error"], + }); break; - case "run_errored": + case "system_item": { + const item = value["data"]; + projection.lines.push(systemItemLine(eventId, item ?? value)); + applyTaskSystemItem(projection, item); + break; + } + case "run_error": projection.lines.push( line( eventId, "error", "Run error", - stringField(entry, "message") ?? "Worker run failed.", + stringField(value, "message") ?? "Worker run failed.", undefined, false, true, ), ); break; - case "extension": - applyExtensionEntry(projection, eventId, entry); - break; default: break; } From f8a7c46cf963680b97ce9f39640a21186f993fd2 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 30 Aug 2026 09:42:46 +0900 Subject: [PATCH 2/4] fix(session): preflight migration and retain timestamps --- crates/protocol/src/lib.rs | 4 + crates/session-store/src/public_snapshot.rs | 45 ++++-- .../session-store/src/worker_session_store.rs | 91 ++++++++++- crates/worker-runtime/src/runtime.rs | 1 + .../src/feature/builtin/memory_extract.rs | 77 +++------- .../src/feature/builtin/worker_observation.rs | 1 + crates/worker/src/session_capture.rs | 143 +++++++++--------- crates/worker/src/spawn/comm_tools.rs | 1 + web/workspace/src/lib/generated/protocol.ts | 6 +- .../src/lib/workspace/console/model.test.ts | 1 + 10 files changed, 223 insertions(+), 147 deletions(-) diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index 339d1d67..bbda812a 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -394,6 +394,8 @@ pub struct SessionSnapshotEntry { /// Stable identity from durable history metadata, or a deterministic /// identity derived from the legacy segment and log position. pub entry_id: String, + /// Timestamp copied from the durable log record that commits this entry. + pub timestamp: u64, pub provenance: SessionEntryProvenance, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub derived_from: Vec, @@ -1536,6 +1538,7 @@ mod tests { session: SessionSnapshot { entries: vec![SessionSnapshotEntry { entry_id: "entry-1".into(), + timestamp: 1, provenance: SessionEntryProvenance::HumanInput, derived_from: Vec::new(), data: SessionSnapshotEntryData::UserInput { @@ -1565,6 +1568,7 @@ mod tests { parsed["data"]["session"]["entries"][0]["kind"], "user_input" ); + assert_eq!(parsed["data"]["session"]["entries"][0]["timestamp"], 1); assert_eq!(parsed["data"]["greeting"]["worker_name"], "test"); assert_eq!(parsed["data"]["greeting"]["tools"][0], "Read"); assert_eq!(parsed["data"]["greeting"]["context_window"], 200_000); diff --git a/crates/session-store/src/public_snapshot.rs b/crates/session-store/src/public_snapshot.rs index d25feb99..07e3a123 100644 --- a/crates/session-store/src/public_snapshot.rs +++ b/crates/session-store/src/public_snapshot.rs @@ -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, history: &[LoggedHistoryEntry], input_segments: Option<&Vec>, + 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 { fn system_entry( item: &SystemItem, entry_id: String, + timestamp: u64, provenance: SessionEntryProvenance, derived_from: Vec, ) -> 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 diff --git a/crates/session-store/src/worker_session_store.rs b/crates/session-store/src/worker_session_store.rs index ddce9e1b..3fa0bae6 100644 --- a/crates/session-store/src/worker_session_store.rs +++ b/crates/session-store/src/worker_session_store.rs @@ -381,6 +381,15 @@ fn segment_log_paths(root: &Path) -> Result, StoreErro } fn migrate_segment_logs_to_v3(root: &Path, session_id: SessionId) -> Result<(), StoreError> { + struct MigrationPlan { + path: PathBuf, + source: Vec, + output: Vec, + } + + // 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 = 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![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(); diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index a1d2a99d..7ac91a3f 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -4588,6 +4588,7 @@ mod tests { session: protocol::SessionSnapshot { entries: vec![protocol::SessionSnapshotEntry { entry_id: "restored-log-entry".to_owned(), + timestamp: 1, provenance: protocol::SessionEntryProvenance::LegacyUnknown, derived_from: Vec::new(), data: protocol::SessionSnapshotEntryData::RunError { diff --git a/crates/worker/src/feature/builtin/memory_extract.rs b/crates/worker/src/feature/builtin/memory_extract.rs index 165efe6a..023dfab0 100644 --- a/crates/worker/src/feature/builtin/memory_extract.rs +++ b/crates/worker/src/feature/builtin/memory_extract.rs @@ -185,12 +185,9 @@ impl Tool for StageMemoryCandidateTool { })?); } if matches!(params.kind, CandidateKind::Preference) - && entries.iter().any(|entry| { - !matches!( - entry.origin, - crate::WorkerHistoryProvenance::HumanInput { .. } - ) - }) + && entries + .iter() + .any(|entry| !matches!(entry.origin, protocol::SessionEntryProvenance::HumanInput)) { return Err(ToolError::InvalidArgument( "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 { - use crate::WorkerHistoryProvenance as Origin; - let mut evidence = EvidenceOrigin { - kind: EvidenceOriginKind::LegacyUnknown, +fn evidence_origin(origin: &protocol::SessionEntryProvenance) -> EvidenceOrigin { + use protocol::SessionEntryProvenance as Origin; + let kind = match origin { + 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, workspace_id: None, runtime_id: None, @@ -335,46 +345,7 @@ fn evidence_origin(origin: &crate::WorkerHistoryProvenance) -> EvidenceOrigin { flow_selector: None, flow_definition_id: 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 { @@ -502,12 +473,10 @@ mod tests { } #[test] - fn human_origin_projects_account_authority_into_evidence() { - let origin = evidence_origin(&crate::WorkerHistoryProvenance::HumanInput { - account_id: "account-1".into(), - }); + fn public_human_origin_preserves_class_without_inventing_account_authority() { + let origin = evidence_origin(&protocol::SessionEntryProvenance::HumanInput); assert_eq!(origin.kind, EvidenceOriginKind::HumanInput); - assert_eq!(origin.account_id.as_deref(), Some("account-1")); + assert_eq!(origin.account_id, None); } #[test] diff --git a/crates/worker/src/feature/builtin/worker_observation.rs b/crates/worker/src/feature/builtin/worker_observation.rs index ad59bb3f..2dc059f6 100644 --- a/crates/worker/src/feature/builtin/worker_observation.rs +++ b/crates/worker/src/feature/builtin/worker_observation.rs @@ -794,6 +794,7 @@ mod tests { }; Some(protocol::SessionSnapshotEntry { entry_id: format!("fake-{index:08}"), + timestamp: index as u64, provenance: protocol::SessionEntryProvenance::LegacyUnknown, derived_from: Vec::new(), data, diff --git a/crates/worker/src/session_capture.rs b/crates/worker/src/session_capture.rs index c0be48d0..19026fb8 100644 --- a/crates/worker/src/session_capture.rs +++ b/crates/worker/src/session_capture.rs @@ -109,7 +109,7 @@ impl ToolPart { #[derive(Debug, Clone)] pub(crate) struct OverviewItem { pub id: SessionEntryRef, - pub origin: WorkerHistoryProvenance, + pub origin: SessionEntryProvenance, pub entry_range: [u64; 2], pub kind: ReferenceKind, pub label: String, @@ -120,7 +120,7 @@ pub(crate) struct OverviewItem { #[derive(Debug, Clone)] pub(crate) struct ReferenceEntry { pub id: SessionEntryRef, - pub origin: WorkerHistoryProvenance, + pub origin: SessionEntryProvenance, pub entry_range: [u64; 2], pub kind: ReferenceKind, pub tool_part: Option, @@ -146,7 +146,7 @@ pub(crate) struct SearchOptions { #[derive(Debug, Clone)] pub(crate) struct SearchHit { pub id: SessionEntryRef, - pub origin: WorkerHistoryProvenance, + pub origin: SessionEntryProvenance, pub kind: ReferenceKind, pub tool_part: Option, pub tool_name: Option, @@ -192,7 +192,7 @@ impl Default for ReadOptions { #[derive(Debug, Clone)] pub(crate) struct ReadEntry { pub id: SessionEntryRef, - pub origin: WorkerHistoryProvenance, + pub origin: SessionEntryProvenance, pub kind: ReferenceKind, pub tool_part: Option, pub tool_name: Option, @@ -211,7 +211,7 @@ pub(crate) struct ReadResult { pub(crate) struct SessionEntryEvidence { pub segment_id: String, pub entry_ref: SessionEntryRef, - pub origin: WorkerHistoryProvenance, + pub origin: SessionEntryProvenance, pub entry_range: [u64; 2], pub kind: ReferenceKind, pub tool_part: Option, @@ -220,10 +220,17 @@ pub(crate) struct SessionEntryEvidence { pub excerpt: String, } +#[derive(Debug, Clone)] +struct CapturedHistoryEntry { + item: Item, + entry_id: session_store::LoggedSessionHistoryEntryId, + origin: SessionEntryProvenance, +} + #[derive(Debug, Clone)] pub(crate) struct SessionCapture { segment_id: String, - entries: Arc>>, + entries: Arc>, overview: Vec, index: Vec, } @@ -280,35 +287,49 @@ impl SessionCapture { SessionSnapshotEntryData::SystemItem { .. } | SessionSnapshotEntryData::RunError { .. } => return None, }; - Some(HistoryEntry::new( + Some(CapturedHistoryEntry { item, - public_snapshot_metadata(entry.entry_id, entry.provenance), - )) + entry_id: session_store::LoggedSessionHistoryEntryId(entry.entry_id), + origin: entry.provenance, + }) }) .collect(); - Self::from_history_entries(segment_id, entries) + Self::from_captured_entries(segment_id, entries) } pub(crate) fn new(segment_id: impl Into, items: Vec) -> Self { let entries = items .into_iter() .enumerate() - .map(|(index, item)| { - let mut metadata = SessionHistoryMetadata::legacy_unknown(); - metadata.entry_id = - session_store::LoggedSessionHistoryEntryId(format!("{index:08}")); - HistoryEntry::new(item, metadata) + .map(|(index, item)| CapturedHistoryEntry { + item, + entry_id: session_store::LoggedSessionHistoryEntryId(format!("{index:08}")), + origin: SessionEntryProvenance::LegacyUnknown, }) .collect(); - Self::from_history_entries(segment_id, entries) + Self::from_captured_entries(segment_id, entries) } pub(crate) fn from_history_entries( segment_id: impl Into, entries: Vec>, + ) -> 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, + entries: Vec, ) -> Self { let segment_id = segment_id.into(); - let entries = Arc::new(entries); let mut overview = Vec::new(); let mut index = Vec::new(); @@ -317,7 +338,7 @@ impl SessionCapture { let entry_range = [idx as u64, idx as u64]; match item { 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; }; let text = content @@ -327,10 +348,10 @@ impl SessionCapture { .join(""); let label = format!("{} message", kind.as_str()); 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 { id: id.clone(), - origin: entry.annotation.origin.clone(), + origin: entry.origin.clone(), entry_range, kind, tool_part: None, @@ -342,7 +363,7 @@ impl SessionCapture { if matches!(kind, ReferenceKind::User | ReferenceKind::Assistant) { overview.push(OverviewItem { id: id.clone(), - origin: entry.annotation.origin.clone(), + origin: entry.origin.clone(), entry_range, kind, label, @@ -356,8 +377,8 @@ impl SessionCapture { } => { let text = format!("{name}\n{arguments}"); index.push(ReferenceEntry { - id: SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id), - origin: entry.annotation.origin.clone(), + id: SessionEntryRef::from_history_entry_id(&entry.entry_id), + origin: entry.origin.clone(), entry_range, kind: ReferenceKind::Tool, tool_part: Some(ToolPart::Input), @@ -383,8 +404,8 @@ impl SessionCapture { content.as_deref().unwrap_or_default(), ); index.push(ReferenceEntry { - id: SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id), - origin: entry.annotation.origin.clone(), + id: SessionEntryRef::from_history_entry_id(&entry.entry_id), + origin: entry.origin.clone(), entry_range, kind: ReferenceKind::Tool, tool_part: Some(ToolPart::Output), @@ -424,7 +445,7 @@ impl SessionCapture { Self { segment_id, - entries, + entries: Arc::new(entries), overview, index, } @@ -607,65 +628,41 @@ impl SessionCapture { } } -fn public_snapshot_metadata( - entry_id: String, - provenance: SessionEntryProvenance, -) -> SessionHistoryMetadata { - let worker = session_store::LoggedWorkerSubject { - workspace_id: None, - runtime_id: None, - 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 } +fn public_provenance(origin: &WorkerHistoryProvenance) -> SessionEntryProvenance { + match origin { + WorkerHistoryProvenance::HumanInput { .. } => SessionEntryProvenance::HumanInput, + WorkerHistoryProvenance::WorkerInput { .. } => SessionEntryProvenance::WorkerInput, + WorkerHistoryProvenance::FlowInstruction { .. } => SessionEntryProvenance::FlowInstruction, + WorkerHistoryProvenance::BackendInstruction { .. } => { + SessionEntryProvenance::BackendInstruction } - SessionEntryProvenance::ModelOutput => WorkerHistoryProvenance::ModelOutput { - worker: worker.clone(), - }, - SessionEntryProvenance::ToolOutput => WorkerHistoryProvenance::ToolOutput { worker }, - SessionEntryProvenance::DerivedSummary => WorkerHistoryProvenance::DerivedSummary, - SessionEntryProvenance::LegacyUnknown => WorkerHistoryProvenance::LegacyUnknown, - }; - SessionHistoryMetadata { - entry_id: session_store::LoggedSessionHistoryEntryId(entry_id), - origin, - derivation: None, + WorkerHistoryProvenance::ModelOutput { .. } => SessionEntryProvenance::ModelOutput, + WorkerHistoryProvenance::ToolOutput { .. } => SessionEntryProvenance::ToolOutput, + WorkerHistoryProvenance::DerivedSummary => SessionEntryProvenance::DerivedSummary, + WorkerHistoryProvenance::LegacyUnknown => SessionEntryProvenance::LegacyUnknown, } } fn message_reference_kind( - origin: &WorkerHistoryProvenance, + origin: &SessionEntryProvenance, provider_role: &Role, ) -> Option { match origin { - WorkerHistoryProvenance::HumanInput { .. } - | WorkerHistoryProvenance::WorkerInput { .. } => Some(ReferenceKind::User), - WorkerHistoryProvenance::ModelOutput { .. } => Some(ReferenceKind::Assistant), - WorkerHistoryProvenance::ToolOutput { .. } => Some(ReferenceKind::Tool), - WorkerHistoryProvenance::LegacyUnknown => match provider_role { + SessionEntryProvenance::HumanInput | SessionEntryProvenance::WorkerInput => { + Some(ReferenceKind::User) + } + SessionEntryProvenance::ModelOutput => Some(ReferenceKind::Assistant), + SessionEntryProvenance::ToolOutput => Some(ReferenceKind::Tool), + SessionEntryProvenance::LegacyUnknown => match provider_role { Role::User => Some(ReferenceKind::User), Role::Assistant => Some(ReferenceKind::Assistant), Role::System => None, }, // Flow/backend/system content remains out of the observation surface // even when represented with a provider user/system role. - WorkerHistoryProvenance::FlowInstruction { .. } - | WorkerHistoryProvenance::BackendInstruction { .. } - | WorkerHistoryProvenance::DerivedSummary => None, + SessionEntryProvenance::FlowInstruction + | SessionEntryProvenance::BackendInstruction + | SessionEntryProvenance::DerivedSummary => None, } } @@ -762,13 +759,13 @@ mod tests { assert_eq!(overview.len(), 1); assert!(matches!( overview[0].origin, - WorkerHistoryProvenance::HumanInput { .. } + SessionEntryProvenance::HumanInput )); let evidence = capture.evidence_for(overview[0].id.as_str()).unwrap(); assert!(evidence.excerpt.ends_with("remember my preference")); assert!(matches!( evidence.origin, - WorkerHistoryProvenance::HumanInput { .. } + SessionEntryProvenance::HumanInput )); } diff --git a/crates/worker/src/spawn/comm_tools.rs b/crates/worker/src/spawn/comm_tools.rs index 9d30d89b..49f9ab78 100644 --- a/crates/worker/src/spawn/comm_tools.rs +++ b/crates/worker/src/spawn/comm_tools.rs @@ -284,6 +284,7 @@ mod tests { .enumerate() .map(|(index, value)| protocol::SessionSnapshotEntry { entry_id: format!("test-{index}"), + timestamp: index as u64, provenance: protocol::SessionEntryProvenance::LegacyUnknown, derived_from: Vec::new(), data: protocol::SessionSnapshotEntryData::RunError { diff --git a/web/workspace/src/lib/generated/protocol.ts b/web/workspace/src/lib/generated/protocol.ts index 39ab33c2..d25d2eb3 100644 --- a/web/workspace/src/lib/generated/protocol.ts +++ b/web/workspace/src/lib/generated/protocol.ts @@ -93,7 +93,11 @@ export type SessionSnapshotEntry = { * Stable identity from durable history metadata, or a deterministic * identity derived from the legacy segment and log position. */ -entry_id: string, provenance: SessionEntryProvenance, derived_from?: Array, } & ({ "kind": "user_input", segments: Array, } | { "kind": "message", role: SessionMessageRole, content: Array, } | { "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, } | { "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, } & ({ "kind": "user_input", segments: Array, } | { "kind": "message", role: SessionMessageRole, content: Array, } | { "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, } | { "kind": "system_item", item_kind: string, content: string, data?: unknown, } | { "kind": "run_error", message: string, }); export type SessionSnapshot = { entries: Array, }; diff --git a/web/workspace/src/lib/workspace/console/model.test.ts b/web/workspace/src/lib/workspace/console/model.test.ts index abc5c350..94c2c018 100644 --- a/web/workspace/src/lib/workspace/console/model.test.ts +++ b/web/workspace/src/lib/workspace/console/model.test.ts @@ -2108,6 +2108,7 @@ Deno.test("snapshot restores TaskStore state from system history", () => { event.data.session = { entries: [{ entry_id: "task-reminder-1", + timestamp: 1, provenance: "backend_instruction", kind: "system_item", item_kind: "task_reminder", From 8493472983274dd56f75f5e44baf2f1493325da3 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 30 Aug 2026 12:18:44 +0900 Subject: [PATCH 3/4] refactor: require annotated session log history --- crates/protocol/src/lib.rs | 4 +- crates/session-store/src/history.rs | 55 +-- .../session-store/src/legacy_session_log.rs | 188 +++++++++ crates/session-store/src/lib.rs | 7 +- crates/session-store/src/public_snapshot.rs | 60 +-- crates/session-store/src/segment.rs | 84 ++-- crates/session-store/src/segment_log.rs | 219 ++++------- crates/session-store/src/system_item.rs | 2 +- .../session-store/src/worker_session_store.rs | 368 +++++++++--------- crates/session-store/tests/fs_store_test.rs | 40 +- crates/session-store/tests/session_test.rs | 93 +++-- crates/tui/src/app.rs | 82 ++-- crates/tui/src/worker_list.rs | 16 +- crates/worker-runtime/src/worker_backend.rs | 9 +- crates/worker/src/controller.rs | 8 +- .../src/feature/builtin/flow_transition.rs | 5 +- crates/worker/src/hook.rs | 2 +- crates/worker/src/in_flight.rs | 10 +- crates/worker/src/ipc/interceptor.rs | 7 +- crates/worker/src/ipc/notify_buffer.rs | 2 +- crates/worker/src/ipc/protocol_session.rs | 15 +- crates/worker/src/ipc/server.rs | 2 +- crates/worker/src/segment_log_sink.rs | 54 +-- crates/worker/src/session_history.rs | 64 ++- crates/worker/src/spawn/registry.rs | 43 +- crates/worker/src/worker.rs | 93 +++-- crates/worker/tests/compact_events_test.rs | 38 +- crates/worker/tests/controller_test.rs | 18 +- 28 files changed, 872 insertions(+), 716 deletions(-) create mode 100644 crates/session-store/src/legacy_session_log.rs diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index bbda812a..8c498f09 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -470,7 +470,7 @@ pub struct SessionToolAttachment { #[serde(tag = "event", content = "data", rename_all = "snake_case")] pub enum Event { /// A user input message was accepted, persisted as - /// `LogEntry::UserInput`, and is about to start a new turn. + /// `LogEntry::AnnotatedUserInput`, and is about to start a new turn. /// Broadcast to every subscribed client so TUI / GUI instances show /// the same user line that reconnect snapshots would replay from /// history; clients must not synthesize a separate pending/fake @@ -491,7 +491,7 @@ pub enum Event { /// of parsing free-text prefixes like `[Notification] …` or /// `[File: …]`. /// - /// One event per `LogEntry::SystemItem` commit. Disk-side and + /// One event per `LogEntry::AnnotatedSystemItem` commit. Disk-side and /// wire-side are 1:1. SystemItem { #[cfg_attr(feature = "typescript", ts(type = "unknown"))] diff --git a/crates/session-store/src/history.rs b/crates/session-store/src/history.rs index 9558f343..e960bfec 100644 --- a/crates/session-store/src/history.rs +++ b/crates/session-store/src/history.rs @@ -1,9 +1,8 @@ //! Serializable history entries with restore-authoritative logical identity and origin. -use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use serde::{Deserialize, Serialize}; -use crate::{LoggedItem, SessionId}; +use crate::LoggedItem; /// Stable logical identity of one model-visible history entry. /// @@ -143,12 +142,15 @@ mod tests { #[test] fn annotated_segment_start_is_restore_visible_without_projecting_metadata() { let session_id = uuid::Uuid::now_v7(); - let history_entry = legacy_logged_history(LoggedItem::Message { - role: LoggedRole::Assistant, - content: vec![crate::LoggedContentPart::Text { - text: "answer".into(), - }], - }); + let history_entry = LoggedHistoryEntry { + item: LoggedItem::Message { + role: LoggedRole::Assistant, + content: vec![crate::LoggedContentPart::Text { + text: "answer".into(), + }], + }, + metadata: LoggedSessionHistoryMetadata::legacy_unknown(), + }; let state = crate::collect_state(&[crate::LogEntry::AnnotatedSegmentStart { ts: 1, session_id, @@ -161,40 +163,3 @@ mod tests { assert_eq!(state.history[0].as_text(), Some("answer")); } } - -/// Legacy Session Logs did not persist annotations. Decode helpers explicitly -/// create `LegacyUnknown`; they never infer Human/System authority from role or -/// plaintext. -pub fn legacy_logged_history(item: LoggedItem) -> LoggedHistoryEntry { - LoggedHistoryEntry { - item, - metadata: LoggedSessionHistoryMetadata::legacy_unknown(), - } -} - -pub fn legacy_segment_history( - session_id: SessionId, - items: impl IntoIterator, -) -> Vec { - items - .into_iter() - .enumerate() - .map(|(index, item)| LoggedHistoryEntry { - item, - metadata: LoggedSessionHistoryMetadata { - // Legacy logs have no persisted entry id. Derive one solely from - // durable segment content rather than minting a new random value - // on every restore/read. The explicit LegacyUnknown origin keeps - // this compatibility identity from becoming trust authority. - entry_id: { - let mut identity = Vec::with_capacity(24); - identity.extend_from_slice(session_id.as_bytes()); - identity.extend_from_slice(&(index as u64).to_be_bytes()); - LoggedSessionHistoryEntryId(format!("l-{}", URL_SAFE_NO_PAD.encode(identity))) - }, - origin: LoggedSessionHistoryOrigin::LegacyUnknown, - derivation: None, - }, - }) - .collect() -} diff --git a/crates/session-store/src/legacy_session_log.rs b/crates/session-store/src/legacy_session_log.rs new file mode 100644 index 00000000..40dd0fd5 --- /dev/null +++ b/crates/session-store/src/legacy_session_log.rs @@ -0,0 +1,188 @@ +//! Versioned decoder for Session schemas that predate canonical annotated history. +//! +//! These types are intentionally private to `session-store`. Current writers, +//! replay, and public projections use [`crate::LogEntry`] exclusively; only the +//! Worker Session schema migration is allowed to deserialize these shapes. + +use agen::llm_client::types::RequestConfig; +use protocol::Segment; +use serde::Deserialize; + +use crate::{ + LogEntry, LoggedHistoryEntry, LoggedItem, LoggedSessionHistoryEntryId, + LoggedSessionHistoryMetadata, LoggedSessionHistoryOrigin, LoggedSystemHistoryEntry, SegmentId, + SegmentOrigin, SessionExtension, SessionId, SystemItem, +}; + +#[derive(Debug, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +enum LegacyHistoryLogEntry { + SegmentStart { + ts: u64, + session_id: SessionId, + system_prompt: Option, + config: RequestConfig, + history: Vec, + #[serde(default)] + forked_from: Option, + #[serde(default)] + compacted_from: Option, + }, + UserInput { + ts: u64, + segments: Vec, + #[serde(default)] + extensions: Vec, + }, + AssistantItem { + ts: u64, + item: LoggedItem, + }, + ToolResult { + ts: u64, + item: LoggedItem, + }, + SystemItem { + ts: u64, + item: SystemItem, + }, +} + +/// Schema-v1 decoder. Non-history records already had their current shape, so +/// they pass through `LogEntry`; legacy history records are converted below. +#[derive(Debug, Deserialize)] +#[serde(untagged)] +enum LegacySessionLogEntryV1 { + History(LegacyHistoryLogEntry), + Current(LogEntry), +} + +/// Schema v2 retained the v1 history shapes while adding non-history records. +/// Keep a distinct type so supported source versions remain explicit rather +/// than turning migration compatibility into the current `LogEntry` contract. +#[derive(Debug, Deserialize)] +#[serde(untagged)] +enum LegacySessionLogEntryV2 { + History(LegacyHistoryLogEntry), + Current(LogEntry), +} + +pub(crate) fn decode_entry( + schema_version: u32, + line: &str, + session_id: SessionId, + segment_id: SegmentId, + line_index: usize, +) -> Result { + let entry = match schema_version { + 1 => match serde_json::from_str::(line)? { + LegacySessionLogEntryV1::History(entry) => Entry::History(entry), + LegacySessionLogEntryV1::Current(entry) => Entry::Current(entry), + }, + 2 => match serde_json::from_str::(line)? { + LegacySessionLogEntryV2::History(entry) => Entry::History(entry), + LegacySessionLogEntryV2::Current(entry) => Entry::Current(entry), + }, + _ => unreachable!("legacy decoder called for unsupported schema {schema_version}"), + }; + Ok(match entry { + Entry::History(entry) => { + canonicalize_history_entry(session_id, segment_id, line_index, entry) + } + Entry::Current(entry) => entry, + }) +} + +enum Entry { + History(LegacyHistoryLogEntry), + Current(LogEntry), +} + +fn legacy_metadata( + segment_id: SegmentId, + line_index: usize, + item_index: usize, +) -> LoggedSessionHistoryMetadata { + let mut identity = Vec::with_capacity(32); + identity.extend_from_slice(segment_id.as_bytes()); + identity.extend_from_slice(&(line_index as u64).to_be_bytes()); + identity.extend_from_slice(&(item_index as u64).to_be_bytes()); + LoggedSessionHistoryMetadata { + entry_id: LoggedSessionHistoryEntryId(format!( + "l-{}", + base64::Engine::encode(&base64::engine::general_purpose::URL_SAFE_NO_PAD, identity) + )), + origin: LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + } +} + +fn canonicalize_history_entry( + _session_id: SessionId, + segment_id: SegmentId, + line_index: usize, + entry: LegacyHistoryLogEntry, +) -> LogEntry { + match entry { + LegacyHistoryLogEntry::SegmentStart { + ts, + session_id, + system_prompt, + config, + history, + forked_from, + compacted_from, + } => LogEntry::AnnotatedSegmentStart { + ts, + session_id, + system_prompt, + config, + history: history + .into_iter() + .enumerate() + .map(|(item_index, item)| LoggedHistoryEntry { + item, + metadata: legacy_metadata(segment_id, line_index, item_index), + }) + .collect(), + forked_from, + compacted_from, + }, + LegacyHistoryLogEntry::UserInput { + ts, + segments, + extensions, + } => LogEntry::AnnotatedUserInput { + ts, + history: vec![LoggedHistoryEntry { + item: LoggedItem::from(agen::Item::user_message(Segment::flatten_to_text( + &segments, + ))), + metadata: legacy_metadata(segment_id, line_index, 0), + }], + segments, + extensions, + }, + LegacyHistoryLogEntry::AssistantItem { ts, item } => LogEntry::AnnotatedAssistantItem { + ts, + entry: LoggedHistoryEntry { + item, + metadata: legacy_metadata(segment_id, line_index, 0), + }, + }, + LegacyHistoryLogEntry::ToolResult { ts, item } => LogEntry::AnnotatedToolResult { + ts, + entry: LoggedHistoryEntry { + item, + metadata: legacy_metadata(segment_id, line_index, 0), + }, + }, + LegacyHistoryLogEntry::SystemItem { ts, item } => LogEntry::AnnotatedSystemItem { + ts, + entry: LoggedSystemHistoryEntry { + item, + metadata: legacy_metadata(segment_id, line_index, 0), + }, + }, + } +} diff --git a/crates/session-store/src/lib.rs b/crates/session-store/src/lib.rs index 07bc0bc9..411a9aa0 100644 --- a/crates/session-store/src/lib.rs +++ b/crates/session-store/src/lib.rs @@ -26,13 +26,14 @@ //! let (session_id, segment_id) = create_segment(&store, SegmentStartState { //! system_prompt: None, //! config: &config, -//! history: &[], +//! history: Vec::new(), //! })?; //! ``` pub mod event_trace; pub mod fs_store; pub mod history; +mod legacy_session_log; pub mod logged_item; pub mod public_snapshot; pub mod segment; @@ -49,11 +50,11 @@ pub use fs_store::FsStore; pub use history::{ LoggedHistoryDerivation, LoggedHistoryEntry, LoggedSessionHistoryEntryId, LoggedSessionHistoryMetadata, LoggedSessionHistoryOrigin, LoggedSystemHistoryEntry, - LoggedWorkerSubject, legacy_logged_history, legacy_segment_history, + LoggedWorkerSubject, }; pub use logged_item::{LoggedContentPart, LoggedItem, LoggedRole, from_logged, to_logged}; pub use segment::{ - SegmentStartState, append_entry, append_system_item, classify_history_item, + SegmentStartState, append_entry, append_system_item, classify_logged_history_entry, create_compacted_segment, create_segment, create_segment_with_ids, ensure_head_or_fork, fork, fork_at, restore, restore_by_segment, save_config_changed, save_delta, save_extension, save_run_completed, save_run_errored, save_turn_end, save_usage, save_user_input, diff --git a/crates/session-store/src/public_snapshot.rs b/crates/session-store/src/public_snapshot.rs index 07e3a123..3109bbe3 100644 --- a/crates/session-store/src/public_snapshot.rs +++ b/crates/session-store/src/public_snapshot.rs @@ -12,13 +12,12 @@ use crate::{ LoggedSessionHistoryOrigin, SessionId, SystemItem, }; -/// Project a complete current-segment log. A valid segment always starts with -/// one of the two SegmentStart records; malformed partial input uses the nil -/// session only to keep the public failure projection deterministic. +/// Project a complete current-segment log. A valid segment starts with one +/// canonical annotated SegmentStart record; malformed partial input uses the +/// nil session only to keep the public failure projection deterministic. pub fn project_current_session_snapshot(log: &[LogEntry]) -> SessionSnapshot { let session_id = log.iter().find_map(|entry| match entry { - LogEntry::SegmentStart { session_id, .. } - | LogEntry::AnnotatedSegmentStart { session_id, .. } => Some(*session_id), + LogEntry::AnnotatedSegmentStart { session_id, .. } => Some(*session_id), _ => None, }); project_session_snapshot(session_id.unwrap_or_else(SessionId::nil), log) @@ -32,20 +31,6 @@ 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, - .. - } => { - session_key = *session_id; - 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, *ts, data)); - } - } - } LogEntry::AnnotatedSegmentStart { ts, session_id, @@ -56,39 +41,18 @@ pub fn project_session_snapshot(session_id: SessionId, log: &[LogEntry]) -> Sess entries.clear(); extend_history(&mut entries, history, None, *ts); } - LogEntry::UserInput { ts, segments, .. } => entries.push(legacy_entry( - &session_key, - log_index, - 0, - *ts, - SessionSnapshotEntryData::UserInput { - segments: segments.clone(), - }, - )), LogEntry::AnnotatedUserInput { 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, *ts, data)); - } - } LogEntry::AnnotatedAssistantItem { ts, entry } | LogEntry::AnnotatedToolResult { ts, entry } => { if let Some(data) = project_item(&entry.item) { entries.push(history_entry(entry, *ts, data)); } } - LogEntry::SystemItem { ts, item } => entries.push(system_entry( - item, - legacy_entry_id(&session_key, log_index, 0), - *ts, - SessionEntryProvenance::LegacyUnknown, - Vec::new(), - )), LogEntry::AnnotatedSystemItem { ts, entry } => entries.push(system_entry( &entry.item, entry.metadata.entry_id.0.clone(), @@ -332,9 +296,9 @@ mod tests { use crate::{LoggedSessionHistoryEntryId, LoggedSessionHistoryMetadata, LoggedWorkerSubject}; #[test] - fn legacy_projection_is_stable_and_hides_reasoning_and_system_prompts() { + fn current_projection_is_stable_and_hides_reasoning_and_system_prompts() { let session_id = crate::new_session_id(); - let log = vec![LogEntry::SegmentStart { + let log = vec![LogEntry::AnnotatedSegmentStart { ts: 1, session_id, system_prompt: None, @@ -358,7 +322,17 @@ mod tests { text: "visible".into(), }], }, - ], + ] + .into_iter() + .map(|item| LoggedHistoryEntry { + item, + metadata: LoggedSessionHistoryMetadata { + entry_id: LoggedSessionHistoryEntryId::new(), + origin: LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + }, + }) + .collect(), forked_from: None, compacted_from: None, }]; diff --git a/crates/session-store/src/segment.rs b/crates/session-store/src/segment.rs index 5be5b320..017cb970 100644 --- a/crates/session-store/src/segment.rs +++ b/crates/session-store/src/segment.rs @@ -4,11 +4,9 @@ //! The caller (typically Worker) holds the Engine directly and calls these //! functions after state-mutating operations. -use crate::logged_item::{LoggedItem, to_logged}; use crate::segment_log::{self, LogEntry, SegmentOrigin}; use crate::store::{Store, StoreError}; -use crate::system_item::SystemItem; -use crate::{SegmentId, SessionId}; +use crate::{LoggedHistoryEntry, LoggedSystemHistoryEntry, SegmentId, SessionId}; use agen::EngineResult; use agen::llm_client::RequestConfig; use agen::llm_client::types::Item; @@ -18,7 +16,7 @@ use protocol::Segment; pub struct SegmentStartState<'a> { pub system_prompt: Option<&'a str>, pub config: &'a RequestConfig, - pub history: &'a [Item], + pub history: Vec, } /// Create a new session + initial segment, writing the initial @@ -44,12 +42,12 @@ pub fn create_segment_with_ids( segment_id: SegmentId, state: SegmentStartState<'_>, ) -> Result<(), StoreError> { - let entry = LogEntry::SegmentStart { + let entry = LogEntry::AnnotatedSegmentStart { ts: segment_log::now_millis(), session_id, system_prompt: state.system_prompt.map(String::from), config: state.config.clone(), - history: to_logged(state.history), + history: state.history.to_vec(), forked_from: None, compacted_from: None, }; @@ -70,12 +68,12 @@ pub fn create_compacted_segment( source_turn_count: usize, ) -> Result { let segment_id = crate::new_segment_id(); - let entry = LogEntry::SegmentStart { + let entry = LogEntry::AnnotatedSegmentStart { ts: segment_log::now_millis(), session_id: source_session_id, system_prompt: state.system_prompt.map(String::from), config: state.config.clone(), - history: to_logged(state.history), + history: state.history.to_vec(), forked_from: None, compacted_from: Some(SegmentOrigin { segment_id: source_segment_id, @@ -154,12 +152,12 @@ pub fn ensure_head_or_fork( } let source_segment_id = *segment_id; let fork_id = crate::new_segment_id(); - let entry = LogEntry::SegmentStart { + let entry = LogEntry::AnnotatedSegmentStart { ts: segment_log::now_millis(), session_id, system_prompt: state.system_prompt.map(String::from), config: state.config.clone(), - history: to_logged(state.history), + history: state.history.to_vec(), forked_from: Some(SegmentOrigin { segment_id: source_segment_id, at_turn_index, @@ -183,8 +181,9 @@ pub fn save_user_input( session_id: SessionId, segment_id: SegmentId, segments: Vec, + history: Vec, ) -> Result<(), StoreError> { - save_user_input_with_extensions(store, session_id, segment_id, segments, Vec::new()) + save_user_input_with_extensions(store, session_id, segment_id, segments, history, Vec::new()) } /// Atomically persist one typed user submission and Runtime-owned session @@ -194,15 +193,17 @@ pub fn save_user_input_with_extensions( session_id: SessionId, segment_id: SegmentId, segments: Vec, + history: Vec, extensions: Vec, ) -> Result<(), StoreError> { append_entry( store, session_id, segment_id, - LogEntry::UserInput { + LogEntry::AnnotatedUserInput { ts: segment_log::now_millis(), segments, + history, extensions, }, ) @@ -220,64 +221,57 @@ pub fn save_delta( store: &impl Store, session_id: SessionId, segment_id: SegmentId, - new_items: &[Item], + new_items: &[LoggedHistoryEntry], ) -> Result<(), StoreError> { if new_items.is_empty() { return Ok(()); } let ts = segment_log::now_millis(); - for item in new_items { + for entry in new_items { + let item = Item::from(entry.item.clone()); if item.is_user_message() { // Already persisted by save_user_input at submit time. continue; } - let entry = classify_history_item(item, ts); + let entry = classify_logged_history_entry(entry.clone(), ts); append_entry(store, session_id, segment_id, entry)?; } Ok(()) } -/// Map one history item to its singular `LogEntry` form. Used by the -/// fallback `save_delta` path and the controller's worker-callback -/// classifier so write classification lives in one place. -pub fn classify_history_item(item: &Item, ts: u64) -> LogEntry { +/// Map one annotated history entry to its singular `LogEntry` form. Used by +/// the fallback `save_delta` path and the controller's worker-callback +/// classifier so write classification lives in one place without discarding +/// identity or provenance. +/// Map one already-annotated history entry to its singular canonical record +/// without changing its identity or provenance. +pub fn classify_logged_history_entry(entry: LoggedHistoryEntry, ts: u64) -> LogEntry { + let item = Item::from(entry.item.clone()); if item.is_tool_result() { - LogEntry::ToolResult { - ts, - item: LoggedItem::from(item), - } - } else if item.is_assistant_message() || item.is_tool_call() || item.is_reasoning() { - LogEntry::AssistantItem { - ts, - item: LoggedItem::from(item), - } + LogEntry::AnnotatedToolResult { ts, entry } } else { - // Defensive: anything else (future Item kinds) routes through - // AssistantItem rather than getting silently dropped. - LogEntry::AssistantItem { - ts, - item: LoggedItem::from(item), - } + // Assistant messages, tool calls, reasoning, and future non-user + // items all use the assistant-side canonical record. + LogEntry::AnnotatedAssistantItem { ts, entry } } } -/// Append a single typed system item as `LogEntry::SystemItem`. Helper -/// for the Worker-side interceptor commit path; mirrors the per-item -/// commit shape used for assistant / tool result entries. +/// Append one typed system item and its history metadata as a canonical +/// `LogEntry::AnnotatedSystemItem`. pub fn append_system_item( store: &impl Store, session_id: SessionId, segment_id: SegmentId, - item: SystemItem, + entry: LoggedSystemHistoryEntry, ) -> Result<(), StoreError> { append_entry( store, session_id, segment_id, - LogEntry::SystemItem { + LogEntry::AnnotatedSystemItem { ts: segment_log::now_millis(), - item, + entry, }, ) } @@ -430,12 +424,12 @@ pub fn fork( ) -> Result<(SessionId, SegmentId), StoreError> { let session_id = crate::new_session_id(); let fork_id = crate::new_segment_id(); - let entry = LogEntry::SegmentStart { + let entry = LogEntry::AnnotatedSegmentStart { ts: segment_log::now_millis(), session_id, system_prompt: state.system_prompt.map(String::from), config: state.config.clone(), - history: to_logged(state.history), + history: state.history.to_vec(), forked_from: None, compacted_from: None, }; @@ -470,7 +464,7 @@ pub fn fork_at( // segment), before any turn completes. entries .iter() - .position(|e| !matches!(e, LogEntry::SegmentStart { .. })) + .position(|e| !matches!(e, LogEntry::AnnotatedSegmentStart { .. })) .unwrap_or(entries.len()) } else { entries @@ -482,12 +476,12 @@ pub fn fork_at( let state = segment_log::collect_state(&entries[..cut]); let fork_id = crate::new_segment_id(); - let entry = LogEntry::SegmentStart { + let entry = LogEntry::AnnotatedSegmentStart { ts: segment_log::now_millis(), session_id: source_session_id, system_prompt: state.system_prompt, config: state.config, - history: to_logged(&state.history), + history: state.annotated_history, forked_from: Some(SegmentOrigin { segment_id: source_id, at_turn_index, diff --git a/crates/session-store/src/segment_log.rs b/crates/session-store/src/segment_log.rs index b03163d2..73e250ea 100644 --- a/crates/session-store/src/segment_log.rs +++ b/crates/session-store/src/segment_log.rs @@ -16,7 +16,6 @@ use serde::{Deserialize, Serialize}; use crate::history::{LoggedHistoryEntry, LoggedSystemHistoryEntry}; use crate::logged_item::LoggedItem; -use crate::system_item::SystemItem; /// A single segment log entry, serialized as one JSONL line. /// @@ -50,28 +49,7 @@ impl SessionExtension { #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "kind", rename_all = "snake_case")] pub enum LogEntry { - /// Segment start. Always the first entry in a segment log. - /// For forked segments, `history` contains the seed state from the parent. - SegmentStart { - ts: u64, - /// Session this segment belongs to. Compaction / fork inherits - /// the source segment's session_id; only fresh "new conversation" - /// segments mint a new session_id. - session_id: crate::SessionId, - system_prompt: Option, - config: RequestConfig, - history: Vec, - /// Origin: forked from a sibling segment at a specific turn boundary. - /// The referenced segment is guaranteed to share `session_id`. - #[serde(default, skip_serializing_if = "Option::is_none")] - forked_from: Option, - /// Origin: compacted from a sibling segment at a specific turn boundary. - /// The referenced segment is guaranteed to share `session_id`. - #[serde(default, skip_serializing_if = "Option::is_none")] - compacted_from: Option, - }, - - /// Schema-v2 segment seed. Retained entries keep their stable logical + /// Canonical segment seed. Retained entries keep their stable logical /// identity and origin across fork/compaction/restore. AnnotatedSegmentStart { ts: u64, @@ -105,22 +83,7 @@ pub enum LogEntry { /// restore conservatively instead of re-running a dangling tool call. Invoke { ts: u64, trigger: InvokeKind }, - /// User input accepted at submit time. Carries the original typed - /// `Vec` so clients can re-render typed atoms (paste chips, - /// file refs) on segment restore. - /// Replay flattens these into a `Item::user_message` for the worker - /// history; the worker layer never sees segments directly. - UserInput { - ts: u64, - segments: Vec, - /// Typed durable state committed atomically with this input record. - /// Runtime-owned Flow invocation uses this to avoid a Backend-instance - /// commit that can get ahead of Worker history. - #[serde(default, skip_serializing_if = "Vec::is_empty")] - extensions: Vec, - }, - - /// Schema-v2 user submission with its exact model-visible entries. Typed + /// Canonical user submission with its exact model-visible entries. Typed /// Flow instructions and caller-attributed input remain separate entries. AnnotatedUserInput { ts: u64, @@ -130,35 +93,19 @@ pub enum LogEntry { history: Vec, }, - /// Schema-v2 model output and metadata committed as one journal record. + /// Canonical model output and metadata committed as one journal record. AnnotatedAssistantItem { ts: u64, entry: LoggedHistoryEntry }, - /// One assistant-side item appended to history — assistant message, - /// reasoning, or tool call. Singular: one entry per history item so - /// the wire-side `Event::*` lane and on-disk LogEntry stay 1:1. - AssistantItem { ts: u64, item: LoggedItem }, - - /// Schema-v2 tool output and metadata committed as one journal record. + /// Canonical tool output and metadata committed as one journal record. AnnotatedToolResult { ts: u64, entry: LoggedHistoryEntry }, - /// One tool-execution result appended to history. - ToolResult { ts: u64, item: LoggedItem }, - - /// Schema-v2 typed system event and model-visible metadata committed + /// Canonical typed system event and model-visible metadata committed /// together. AnnotatedSystemItem { ts: u64, entry: LoggedSystemHistoryEntry, }, - /// One typed agent-injected system item: notification, child-Worker - /// lifecycle event, `@` / `/` resolution payload. Each - /// `SystemItem` carries kind metadata that the LLM - /// itself never sees (the LLM gets `Item::system_message` with the - /// item's denormalised `body`), but live clients and replay paths - /// dispatch on `kind` for typed rendering. - SystemItem { ts: u64, item: SystemItem }, - /// Turn boundary. Records the turn count after increment. TurnEnd { ts: u64, turn_count: usize }, @@ -260,6 +207,10 @@ pub struct RestoredState { pub system_prompt: Option, pub config: RequestConfig, pub history: Vec, + /// Canonical persisted history with stable identity and provenance. This is + /// the authority for rewrites, forks, and annotated restore; `history` is + /// retained as the model-facing item projection. + pub annotated_history: Vec, pub turn_count: usize, /// AgentTurns consumed by the active paused/yielded logical run. pub active_run_turn_count: Option, @@ -276,7 +227,7 @@ pub struct RestoredState { /// session-store は domain を不透明扱いし、各ドメインが自前で fold する。 pub extensions: Vec<(String, serde_json::Value)>, /// User submissions in original typed form, in submit order. - /// One entry per `LogEntry::UserInput`; the K-th entry corresponds to + /// One entry per `LogEntry::AnnotatedUserInput`; the K-th entry corresponds to /// the K-th `Item::user_message` derived during replay (modulo /// pre-compaction history seeded via `SegmentStart.history`, whose /// original segments are not preserved). Used by clients to re-render @@ -291,6 +242,7 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState { system_prompt: None, config: RequestConfig::default(), history: Vec::new(), + annotated_history: Vec::new(), turn_count: 0, active_run_turn_count: None, last_run_interrupted: false, @@ -304,18 +256,6 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState { state.entries_count += 1; match entry { - LogEntry::SegmentStart { - session_id, - system_prompt, - config, - history, - .. - } => { - state.session_id = Some(*session_id); - state.system_prompt = system_prompt.clone(); - state.config = config.clone(); - state.history = history.iter().cloned().map(Item::from).collect(); - } LogEntry::AnnotatedSegmentStart { session_id, system_prompt, @@ -326,6 +266,7 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState { state.session_id = Some(*session_id); state.system_prompt = system_prompt.clone(); state.config = config.clone(); + state.annotated_history = history.clone(); state.history = history .iter() .cloned() @@ -338,26 +279,13 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState { state.last_run_interrupted = true; state.active_run_turn_count = Some(0); } - LogEntry::UserInput { - segments, - extensions, - .. - } => { - let text = Segment::flatten_to_text(segments); - state.history.push(Item::user_message(text)); - state.user_segments.push(segments.clone()); - state.extensions.extend( - extensions - .iter() - .map(|extension| (extension.domain.clone(), extension.payload.clone())), - ); - } LogEntry::AnnotatedUserInput { segments, extensions, history, .. } => { + state.annotated_history.extend(history.iter().cloned()); state .history .extend(history.iter().cloned().map(|entry| Item::from(entry.item))); @@ -370,20 +298,16 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState { } LogEntry::AnnotatedAssistantItem { entry, .. } | LogEntry::AnnotatedToolResult { entry, .. } => { + state.annotated_history.push(entry.clone()); state.history.push(Item::from(entry.item.clone())); } LogEntry::AnnotatedSystemItem { entry, .. } => { + state.annotated_history.push(LoggedHistoryEntry { + item: LoggedItem::from(entry.item.to_history_item()), + metadata: entry.metadata.clone(), + }); state.history.push(entry.item.to_history_item()); } - LogEntry::AssistantItem { item, .. } => { - state.history.push(Item::from(item.clone())); - } - LogEntry::ToolResult { item, .. } => { - state.history.push(Item::from(item.clone())); - } - LogEntry::SystemItem { item, .. } => { - state.history.push(item.to_history_item()); - } LogEntry::TurnEnd { turn_count, .. } => { if let Some(active_turn_count) = &mut state.active_run_turn_count { *active_turn_count += turn_count.saturating_sub(state.turn_count); @@ -465,6 +389,20 @@ pub fn now_millis() -> u64 { #[cfg(test)] mod tests { use super::*; + use crate::{ + LoggedSessionHistoryEntryId, LoggedSessionHistoryMetadata, LoggedSessionHistoryOrigin, + }; + + fn annotated(item: Item) -> LoggedHistoryEntry { + LoggedHistoryEntry { + item: LoggedItem::from(item), + metadata: LoggedSessionHistoryMetadata { + entry_id: LoggedSessionHistoryEntryId::new(), + origin: LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + }, + } + } #[test] fn replay_empty() { @@ -476,12 +414,12 @@ mod tests { #[test] fn replay_segment_start_sets_initial_state() { - let state = collect_state(&[LogEntry::SegmentStart { + let state = collect_state(&[LogEntry::AnnotatedSegmentStart { ts: 1000, session_id: uuid::Uuid::nil(), system_prompt: Some("You are helpful.".into()), config: RequestConfig::default().with_max_tokens(1024), - history: vec![Item::user_message("seed").into()], + history: vec![annotated(Item::user_message("seed"))], forked_from: None, compacted_from: None, }]); @@ -494,7 +432,7 @@ mod tests { #[test] fn replay_full_turn() { let state = collect_state(&[ - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts: 1000, session_id: uuid::Uuid::nil(), system_prompt: None, @@ -503,14 +441,15 @@ mod tests { forked_from: None, compacted_from: None, }, - LogEntry::UserInput { + LogEntry::AnnotatedUserInput { ts: 2000, extensions: vec![], segments: vec![Segment::text("Hello")], + history: vec![annotated(Item::user_message("Hello"))], }, - LogEntry::AssistantItem { + LogEntry::AnnotatedAssistantItem { ts: 3000, - item: Item::assistant_message("Hi!").into(), + entry: annotated(Item::assistant_message("Hi!")), }, LogEntry::TurnEnd { ts: 3100, @@ -531,7 +470,7 @@ mod tests { #[test] fn replay_incomplete_invoke_is_interrupted() { let state = collect_state(&[ - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts: 1000, session_id: uuid::Uuid::nil(), system_prompt: None, @@ -544,14 +483,15 @@ mod tests { ts: 2000, trigger: InvokeKind::UserSend, }, - LogEntry::UserInput { + LogEntry::AnnotatedUserInput { ts: 2001, extensions: vec![], segments: vec![Segment::text("run a tool")], + history: vec![annotated(Item::user_message("run a tool"))], }, - LogEntry::AssistantItem { + LogEntry::AnnotatedAssistantItem { ts: 3000, - item: Item::tool_call("call_1", "side_effect", "{}").into(), + entry: annotated(Item::tool_call("call_1", "side_effect", "{}")), }, ]); @@ -561,7 +501,7 @@ mod tests { #[test] fn replay_with_tool_calls() { let state = collect_state(&[ - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts: 1000, session_id: uuid::Uuid::nil(), system_prompt: None, @@ -570,22 +510,27 @@ mod tests { forked_from: None, compacted_from: None, }, - LogEntry::UserInput { + LogEntry::AnnotatedUserInput { ts: 2000, extensions: vec![], segments: vec![Segment::text("Check weather")], + history: vec![annotated(Item::user_message("Check weather"))], }, - LogEntry::AssistantItem { + LogEntry::AnnotatedAssistantItem { ts: 3000, - item: Item::tool_call("call_1", "get_weather", r#"{"city":"Tokyo"}"#).into(), + entry: annotated(Item::tool_call( + "call_1", + "get_weather", + r#"{"city":"Tokyo"}"#, + )), }, - LogEntry::ToolResult { + LogEntry::AnnotatedToolResult { ts: 3500, - item: Item::tool_result("call_1", "Sunny, 25C").into(), + entry: annotated(Item::tool_result("call_1", "Sunny, 25C")), }, - LogEntry::AssistantItem { + LogEntry::AnnotatedAssistantItem { ts: 4000, - item: Item::assistant_message("It's sunny in Tokyo!").into(), + entry: annotated(Item::assistant_message("It's sunny in Tokyo!")), }, LogEntry::TurnEnd { ts: 4100, @@ -599,9 +544,9 @@ mod tests { #[test] fn replay_restores_durable_tool_image_detail() { - let entry = LogEntry::ToolResult { + let entry = LogEntry::AnnotatedToolResult { ts: 3500, - item: Item::tool_result_item_with_attachments( + entry: annotated(Item::tool_result_item_with_attachments( "call_image", "attached", None, @@ -609,8 +554,7 @@ mod tests { vec![agen::tool::Attachment::Image( agen::tool::ImageAttachment::new("image/png", b"durable-image".to_vec()), )], - ) - .into(), + )), }; let persisted = serde_json::to_string(&entry).unwrap(); let restored_entry: LogEntry = serde_json::from_str(&persisted).unwrap(); @@ -630,7 +574,7 @@ mod tests { #[test] fn replay_config_changed() { let state = collect_state(&[ - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts: 1000, session_id: uuid::Uuid::nil(), system_prompt: None, @@ -650,7 +594,7 @@ mod tests { #[test] fn replay_llm_usage_appends_to_usage_history() { let state = collect_state(&[ - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts: 1000, session_id: uuid::Uuid::nil(), system_prompt: None, @@ -659,10 +603,11 @@ mod tests { forked_from: None, compacted_from: None, }, - LogEntry::UserInput { + LogEntry::AnnotatedUserInput { ts: 2000, extensions: vec![], segments: vec![Segment::text("hi")], + history: vec![annotated(Item::user_message("hi"))], }, LogEntry::LlmUsage { ts: 2100, @@ -672,9 +617,9 @@ mod tests { cache_write_tokens: 0, output_tokens: 10, }, - LogEntry::AssistantItem { + LogEntry::AnnotatedAssistantItem { ts: 2200, - item: Item::assistant_message("yo").into(), + entry: annotated(Item::assistant_message("yo")), }, LogEntry::LlmUsage { ts: 3100, @@ -698,7 +643,7 @@ mod tests { #[test] fn replay_without_llm_usage_keeps_usage_history_empty() { let state = collect_state(&[ - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts: 1000, session_id: uuid::Uuid::nil(), system_prompt: None, @@ -707,10 +652,11 @@ mod tests { forked_from: None, compacted_from: None, }, - LogEntry::UserInput { + LogEntry::AnnotatedUserInput { ts: 2000, extensions: vec![], segments: vec![Segment::text("hi")], + history: vec![annotated(Item::user_message("hi"))], }, ]); assert!(state.usage_history.is_empty()); @@ -771,7 +717,7 @@ mod tests { #[test] fn replay_invoke_marker_only_mutates_interrupted_state() { let state = collect_state(&[ - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts: 0, session_id: uuid::Uuid::nil(), system_prompt: None, @@ -784,10 +730,11 @@ mod tests { ts: 100, trigger: InvokeKind::UserSend, }, - LogEntry::UserInput { + LogEntry::AnnotatedUserInput { ts: 101, extensions: vec![], segments: vec![Segment::text("hi")], + history: vec![annotated(Item::user_message("hi"))], }, LogEntry::TurnEnd { ts: 200, @@ -806,7 +753,7 @@ mod tests { #[test] fn replay_paused_turn_abandoned_clears_interrupted_marker() { let state = collect_state(&[ - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts: 0, session_id: uuid::Uuid::nil(), system_prompt: None, @@ -830,7 +777,7 @@ mod tests { #[test] fn replay_restores_active_run_budget_across_compaction_checkpoint() { let state = collect_state(&[ - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts: 0, session_id: uuid::Uuid::nil(), system_prompt: None, @@ -861,7 +808,7 @@ mod tests { })) .expect("legacy run-completed entry"); let state = collect_state(&[ - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts: 0, session_id: uuid::Uuid::nil(), system_prompt: None, @@ -924,7 +871,7 @@ mod tests { #[test] fn replay_extension_collects_domain_payload_pairs() { let state = collect_state(&[ - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts: 1000, session_id: uuid::Uuid::nil(), system_prompt: None, @@ -983,9 +930,12 @@ mod tests { #[test] fn user_input_extensions_restore_with_the_same_committed_input() { let segments = vec![Segment::text("Flow instructions"), Segment::text("Ticket")]; - let entry = LogEntry::UserInput { + let entry = LogEntry::AnnotatedUserInput { ts: 9999, segments: segments.clone(), + history: vec![annotated(Item::user_message(Segment::flatten_to_text( + &segments, + )))], extensions: vec![SessionExtension::new( "flow.runtime.v1", serde_json::json!({ "state": "implement", "revision": 0 }), @@ -1000,7 +950,7 @@ mod tests { assert_eq!(state.extensions[0].1["state"], "implement"); } - /// Mixed segments survive a JSON round-trip through `LogEntry::UserInput`, + /// Mixed segments survive a JSON round-trip through `LogEntry::AnnotatedUserInput`, /// and `collect_state` derives `Item::user_message` from the flattened /// text while preserving the original segments separately. This covers /// the segments → flatten → Item replay path from the ticket. @@ -1020,16 +970,19 @@ mod tests { path: "src/main.rs".into(), }, ]; - let entry = LogEntry::UserInput { + let entry = LogEntry::AnnotatedUserInput { ts: 4242, extensions: vec![], segments: segments.clone(), + history: vec![annotated(Item::user_message(Segment::flatten_to_text( + &segments, + )))], }; // JSON round-trip preserves the variant byte-for-byte. let json = serde_json::to_string(&entry).unwrap(); let parsed: LogEntry = serde_json::from_str(&json).unwrap(); let state = collect_state(&[ - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts: 1, session_id: uuid::Uuid::nil(), system_prompt: None, diff --git a/crates/session-store/src/system_item.rs b/crates/session-store/src/system_item.rs index 8ab59b56..cf04db4d 100644 --- a/crates/session-store/src/system_item.rs +++ b/crates/session-store/src/system_item.rs @@ -8,7 +8,7 @@ //! `kind` instead of parsing text prefixes like `[Notification] …` or //! `[File: …]`. //! -//! Persisted as the payload of [`crate::LogEntry::SystemItem`] (one +//! Persisted as the payload of [`crate::LogEntry::AnnotatedSystemItem`] (one //! entry per item), and broadcast live as the payload of //! `Event::SystemItem` on the wire. //! diff --git a/crates/session-store/src/worker_session_store.rs b/crates/session-store/src/worker_session_store.rs index 3fa0bae6..dcc12820 100644 --- a/crates/session-store/src/worker_session_store.rs +++ b/crates/session-store/src/worker_session_store.rs @@ -12,11 +12,7 @@ use crate::event_trace::TraceEntry; use crate::segment_log::LogEntry; use crate::store::{Store, StoreError}; -use crate::{ - LoggedHistoryEntry, LoggedItem, LoggedSessionHistoryEntryId, LoggedSessionHistoryMetadata, - LoggedSessionHistoryOrigin, LoggedSystemHistoryEntry, SegmentId, SessionId, -}; -use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; +use crate::{SegmentId, SessionId}; use serde::{Deserialize, Serialize}; use std::fs::{self, File, OpenOptions}; use std::io::{Read, Seek, SeekFrom, Write}; @@ -56,7 +52,11 @@ impl WorkerSessionStore { validate_canonical_segment_logs(&root)?; } PREVIOUS_SESSION_SCHEMA_VERSION | LEGACY_SESSION_SCHEMA_VERSION => { - migrate_segment_logs_to_v3(&root, manifest.session_id)?; + migrate_segment_logs_to_v3( + &root, + manifest.session_id, + manifest.schema_version, + )?; manifest.schema_version = SESSION_SCHEMA_VERSION; atomic_write_json(&root.join(SESSION_FILE), &manifest)?; } @@ -151,13 +151,7 @@ impl WorkerSessionStore { .join(format!("{segment_id}.trace.jsonl")) } - fn append_log_entry( - &self, - path: &Path, - session_id: SessionId, - segment_id: SegmentId, - entry: &LogEntry, - ) -> Result<(), StoreError> { + fn append_log_entry(&self, path: &Path, entry: &LogEntry) -> Result<(), StoreError> { let _guard = self .append_lock .lock() @@ -172,9 +166,8 @@ impl WorkerSessionStore { file.seek(SeekFrom::Start(0))?; let mut existing = Vec::new(); file.read_to_end(&mut existing)?; - let line_index = parse_jsonl::(&existing)?.len(); - let entry = canonicalize_log_entry(session_id, segment_id, line_index, entry.clone()); - let line = serde_json::to_string(&entry)?; + parse_jsonl::(&existing)?; + let line = serde_json::to_string(entry)?; let mut record = Vec::with_capacity(line.len() + 1); record.extend_from_slice(line.as_bytes()); record.push(b'\n'); @@ -232,7 +225,7 @@ impl Store for WorkerSessionStore { entry: &LogEntry, ) -> Result<(), StoreError> { self.ensure_session(session_id, true)?; - self.append_log_entry(&self.log_path(segment_id), session_id, segment_id, entry) + self.append_log_entry(&self.log_path(segment_id), entry) } fn read_all( @@ -285,9 +278,8 @@ impl Store for WorkerSessionStore { ) -> Result<(), StoreError> { self.ensure_session(session_id, true)?; let mut content = Vec::new(); - for (line_index, entry) in entries.iter().enumerate() { - let entry = canonicalize_log_entry(session_id, segment_id, line_index, entry.clone()); - serde_json::to_writer(&mut content, &entry)?; + for entry in entries { + serde_json::to_writer(&mut content, entry)?; content.push(b'\n'); } atomic_write_bytes(&self.log_path(segment_id), &content)?; @@ -380,7 +372,11 @@ fn segment_log_paths(root: &Path) -> Result, StoreErro Ok(paths) } -fn migrate_segment_logs_to_v3(root: &Path, session_id: SessionId) -> Result<(), StoreError> { +fn migrate_segment_logs_to_v3( + root: &Path, + session_id: SessionId, + source_schema_version: u32, +) -> Result<(), StoreError> { struct MigrationPlan { path: PathBuf, source: Vec, @@ -392,21 +388,14 @@ fn migrate_segment_logs_to_v3(root: &Path, session_id: SessionId) -> Result<(), let mut plans = Vec::new(); for (segment_id, path) in segment_log_paths(root)? { let source = fs::read(&path)?; - let entries: Vec = parse_jsonl(&source).map_err(|error| StoreError::Corrupt { - line: 0, - message: format!( - "cannot migrate Worker Session log {}: {error}", - path.display() - ), - })?; - let canonical = entries - .into_iter() - .enumerate() - .map(|(line_index, entry)| { - canonicalize_log_entry(session_id, segment_id, line_index, entry) - }) - .collect::>(); - validate_canonical_entries(&path, &canonical)?; + let canonical = parse_legacy_jsonl(source_schema_version, session_id, segment_id, &source) + .map_err(|error| StoreError::Corrupt { + line: 0, + message: format!( + "cannot migrate Worker Session log {}: {error}", + path.display() + ), + })?; let mut output = Vec::new(); for entry in canonical { serde_json::to_writer(&mut output, &entry)?; @@ -442,120 +431,33 @@ fn migrate_segment_logs_to_v3(root: &Path, session_id: SessionId) -> Result<(), fn validate_canonical_segment_logs(root: &Path) -> Result<(), StoreError> { for (_, path) in segment_log_paths(root)? { - let entries: Vec = parse_jsonl(&fs::read(&path)?)?; - validate_canonical_entries(&path, &entries)?; + let _: Vec = parse_jsonl(&fs::read(&path)?)?; } Ok(()) } -fn validate_canonical_entries(path: &Path, entries: &[LogEntry]) -> Result<(), StoreError> { - for (line_index, entry) in entries.iter().enumerate() { - if matches!( - entry, - LogEntry::SegmentStart { .. } - | LogEntry::UserInput { .. } - | LogEntry::AssistantItem { .. } - | LogEntry::ToolResult { .. } - | LogEntry::SystemItem { .. } - ) { - return Err(StoreError::Corrupt { - line: line_index + 1, - message: format!( - "Worker Session schema v3 contains legacy history record in {}", - path.display() - ), - }); - } - } - Ok(()) -} - -fn legacy_metadata( - _session_id: SessionId, - segment_id: SegmentId, - line_index: usize, - item_index: usize, -) -> LoggedSessionHistoryMetadata { - let mut identity = Vec::with_capacity(32); - identity.extend_from_slice(segment_id.as_bytes()); - identity.extend_from_slice(&(line_index as u64).to_be_bytes()); - identity.extend_from_slice(&(item_index as u64).to_be_bytes()); - LoggedSessionHistoryMetadata { - entry_id: LoggedSessionHistoryEntryId(format!("l-{}", URL_SAFE_NO_PAD.encode(identity))), - origin: LoggedSessionHistoryOrigin::LegacyUnknown, - derivation: None, - } -} - -fn canonicalize_log_entry( +fn parse_legacy_jsonl( + schema_version: u32, session_id: SessionId, segment_id: SegmentId, - line_index: usize, - entry: LogEntry, -) -> LogEntry { - match entry { - LogEntry::SegmentStart { - ts, - session_id, - system_prompt, - config, - history, - forked_from, - compacted_from, - } => LogEntry::AnnotatedSegmentStart { - ts, - session_id, - system_prompt, - config, - history: history - .into_iter() - .enumerate() - .map(|(item_index, item)| LoggedHistoryEntry { - item, - metadata: legacy_metadata(session_id, segment_id, line_index, item_index), - }) - .collect(), - forked_from, - compacted_from, - }, - LogEntry::UserInput { - ts, - segments, - extensions, - } => LogEntry::AnnotatedUserInput { - ts, - history: vec![LoggedHistoryEntry { - item: LoggedItem::from(agen::Item::user_message( - protocol::Segment::flatten_to_text(&segments), - )), - metadata: legacy_metadata(session_id, segment_id, line_index, 0), - }], - segments, - extensions, - }, - LogEntry::AssistantItem { ts, item } => LogEntry::AnnotatedAssistantItem { - ts, - entry: LoggedHistoryEntry { - item, - metadata: legacy_metadata(session_id, segment_id, line_index, 0), - }, - }, - LogEntry::ToolResult { ts, item } => LogEntry::AnnotatedToolResult { - ts, - entry: LoggedHistoryEntry { - item, - metadata: legacy_metadata(session_id, segment_id, line_index, 0), - }, - }, - LogEntry::SystemItem { ts, item } => LogEntry::AnnotatedSystemItem { - ts, - entry: LoggedSystemHistoryEntry { - item, - metadata: legacy_metadata(session_id, segment_id, line_index, 0), - }, - }, - canonical => canonical, - } + bytes: &[u8], +) -> Result, serde_json::Error> { + let text = std::str::from_utf8(bytes).map_err(|error| { + serde_json::Error::io(std::io::Error::new(std::io::ErrorKind::InvalidData, error)) + })?; + text.lines() + .enumerate() + .filter(|(_, line)| !line.trim().is_empty()) + .map(|(line_index, line)| { + crate::legacy_session_log::decode_entry( + schema_version, + line, + session_id, + segment_id, + line_index, + ) + }) + .collect() } fn atomic_write_json(path: &Path, value: &T) -> Result<(), StoreError> { @@ -659,7 +561,21 @@ fn truncate_uncommitted_tail(file: &mut File) -> std::io::Result { #[cfg(test)] mod tests { use super::*; - use crate::{Store, new_segment_id, new_session_id}; + use crate::{ + LoggedHistoryEntry, LoggedItem, LoggedSessionHistoryEntryId, LoggedSessionHistoryMetadata, + LoggedSessionHistoryOrigin, Store, new_segment_id, new_session_id, + }; + + fn annotated(item: agen::Item) -> LoggedHistoryEntry { + LoggedHistoryEntry { + item: LoggedItem::from(item), + metadata: LoggedSessionHistoryMetadata { + entry_id: LoggedSessionHistoryEntryId::new(), + origin: LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + }, + } + } #[test] fn canonical_layout_and_single_session_invariant() { @@ -748,26 +664,27 @@ mod tests { ) .unwrap(); let source = vec![ - 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, - }, - LogEntry::UserInput { - ts: 2, - segments: vec![protocol::Segment::Text { - content: "hello".into(), - }], - extensions: Vec::new(), - }, - LogEntry::AssistantItem { - ts: 3, - item: LoggedItem::from(agen::Item::assistant_message("reply")), - }, + serde_json::json!({ + "kind": "segment_start", + "ts": 1, + "session_id": session_id, + "system_prompt": null, + "config": agen::llm_client::RequestConfig::default(), + "history": [LoggedItem::from(agen::Item::assistant_message("prior"))], + "forked_from": null, + "compacted_from": null + }), + serde_json::json!({ + "kind": "user_input", + "ts": 2, + "segments": [{ "kind": "text", "content": "hello" }], + "extensions": [] + }), + serde_json::json!({ + "kind": "assistant_item", + "ts": 3, + "item": LoggedItem::from(agen::Item::assistant_message("reply")) + }), ]; let path = root .path() @@ -829,15 +746,16 @@ mod tests { .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 valid_entry = serde_json::json!({ + "kind": "segment_start", + "ts": 1, + "session_id": session_id, + "system_prompt": null, + "config": agen::llm_client::RequestConfig::default(), + "history": [LoggedItem::from(agen::Item::assistant_message("prior"))], + "forked_from": null, + "compacted_from": null + }); let mut valid_bytes = serde_json::to_vec(&valid_entry).unwrap(); valid_bytes.push(b'\n'); fs::write(&valid_path, &valid_bytes).unwrap(); @@ -862,7 +780,7 @@ mod tests { } #[test] - fn schema_v3_rejects_legacy_records_and_new_writes_are_canonical() { + fn current_jsonl_requires_annotations_across_append_rewrite_and_reopen() { let root = tempfile::tempdir().unwrap(); let session_id = new_session_id(); let segment_id = new_segment_id(); @@ -871,12 +789,12 @@ mod tests { .create_segment( session_id, segment_id, - &[LogEntry::SegmentStart { + &[LogEntry::AnnotatedSegmentStart { ts: 1, session_id, system_prompt: None, config: agen::llm_client::RequestConfig::default(), - history: Vec::new(), + history: vec![annotated(agen::Item::user_message("seed"))], forked_from: None, compacted_from: None, }], @@ -886,11 +804,94 @@ mod tests { .append( session_id, segment_id, - &LogEntry::UserInput { + &LogEntry::AnnotatedAssistantItem { + ts: 2, + entry: annotated(agen::Item::assistant_message("reply")), + }, + ) + .unwrap(); + + let before_rewrite = store.read_all(session_id, segment_id).unwrap(); + store + .create_segment(session_id, segment_id, &before_rewrite) + .unwrap(); + drop(store); + + let reopened = WorkerSessionStore::new(root.path()).unwrap(); + let restored = reopened.read_all(session_id, segment_id).unwrap(); + assert_eq!( + serde_json::to_value(&restored).unwrap(), + serde_json::to_value(&before_rewrite).unwrap() + ); + for entry in &restored { + match entry { + LogEntry::AnnotatedSegmentStart { history, .. } => assert!(history.iter().all( + |entry| !entry.metadata.entry_id.0.is_empty() + && matches!( + entry.metadata.origin, + LoggedSessionHistoryOrigin::LegacyUnknown + ) + )), + LogEntry::AnnotatedAssistantItem { entry, .. } => { + assert!(!entry.metadata.entry_id.0.is_empty()); + assert!(matches!( + entry.metadata.origin, + LoggedSessionHistoryOrigin::LegacyUnknown + )); + } + _ => {} + } + } + + let log = fs::read_to_string(reopened.log_path(segment_id)).unwrap(); + for line in log.lines() { + let value: serde_json::Value = serde_json::from_str(line).unwrap(); + let kind = value["kind"].as_str().unwrap(); + assert!( + !matches!( + kind, + "segment_start" + | "user_input" + | "assistant_item" + | "tool_result" + | "system_item" + ), + "current-schema JSONL contains legacy history record: {kind}" + ); + } + } + + #[test] + fn schema_v3_rejects_legacy_records_and_new_writes_are_canonical() { + let root = tempfile::tempdir().unwrap(); + let session_id = new_session_id(); + let segment_id = new_segment_id(); + let store = WorkerSessionStore::new(root.path()).unwrap(); + store + .create_segment( + session_id, + segment_id, + &[LogEntry::AnnotatedSegmentStart { + ts: 1, + session_id, + system_prompt: None, + config: agen::llm_client::RequestConfig::default(), + history: vec![annotated(agen::Item::assistant_message("seed"))], + forked_from: None, + compacted_from: None, + }], + ) + .unwrap(); + store + .append( + session_id, + segment_id, + &LogEntry::AnnotatedUserInput { ts: 2, segments: vec![protocol::Segment::Text { content: "new".into(), }], + history: vec![annotated(agen::Item::user_message("new"))], extensions: Vec::new(), }, ) @@ -907,12 +908,11 @@ mod tests { let mut file = OpenOptions::new().append(true).open(path).unwrap(); serde_json::to_writer( &mut file, - &LogEntry::SystemItem { - ts: 3, - item: crate::SystemItem::LegacyIgnored { - slug: "legacy".into(), - }, - }, + &serde_json::json!({ + "kind": "system_item", + "ts": 3, + "item": { "kind": "legacy_ignored", "slug": "legacy" } + }), ) .unwrap(); file.write_all(b"\n").unwrap(); diff --git a/crates/session-store/tests/fs_store_test.rs b/crates/session-store/tests/fs_store_test.rs index 63028d29..430a685b 100644 --- a/crates/session-store/tests/fs_store_test.rs +++ b/crates/session-store/tests/fs_store_test.rs @@ -1,12 +1,25 @@ use agen::EngineResult; use agen::llm_client::types::{Item, RequestConfig}; use session_store::{ - FsStore, LogEntry, Store, TraceEntry, collect_state, new_segment_id, new_session_id, + 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 { + entry_id: LoggedSessionHistoryEntryId::new(), + origin: LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + }, + } +} + fn nil_session_start(ts: u64, session_id: uuid::Uuid) -> LogEntry { - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts, session_id, system_prompt: None, @@ -25,7 +38,7 @@ fn round_trip_write_and_read() { let segid = new_segment_id(); let entries = vec![ - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts: 1000, session_id: sid, system_prompt: Some("You are helpful.".into()), @@ -34,14 +47,15 @@ fn round_trip_write_and_read() { forked_from: None, compacted_from: None, }, - LogEntry::UserInput { + LogEntry::AnnotatedUserInput { ts: 2000, extensions: vec![], segments: vec![protocol::Segment::text("Hello")], + history: vec![annotated(Item::user_message("Hello"))], }, - LogEntry::AssistantItem { + LogEntry::AnnotatedAssistantItem { ts: 3000, - item: Item::assistant_message("Hi there!").into(), + entry: annotated(Item::assistant_message("Hi there!")), }, LogEntry::TurnEnd { ts: 3100, @@ -79,14 +93,14 @@ fn create_segment_writes_all_entries() { let sid = new_session_id(); let segid = new_segment_id(); - let entries = [LogEntry::SegmentStart { + let entries = [LogEntry::AnnotatedSegmentStart { ts: 1000, session_id: sid, system_prompt: None, config: RequestConfig::default(), history: vec![ - Item::user_message("seed").into(), - Item::assistant_message("ok").into(), + annotated(Item::user_message("seed")), + annotated(Item::assistant_message("ok")), ], forked_from: None, compacted_from: None, @@ -205,7 +219,7 @@ fn read_entry_count_matches_append_tally() { let segid = new_segment_id(); let entries = [ - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts: 1000, session_id: sid, system_prompt: None, @@ -214,10 +228,11 @@ fn read_entry_count_matches_append_tally() { forked_from: None, compacted_from: None, }, - LogEntry::UserInput { + LogEntry::AnnotatedUserInput { ts: 2000, extensions: vec![], segments: vec![protocol::Segment::text("Hello")], + history: vec![annotated(Item::user_message("Hello"))], }, ]; @@ -254,10 +269,11 @@ fn unterminated_utf8_tail_is_ignored_and_replaced_on_append() { assert_eq!(store.read_all(sid, segid).unwrap().len(), 1); assert_eq!(store.read_entry_count(sid, segid).unwrap(), 1); - let next = LogEntry::UserInput { + 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(); diff --git a/crates/session-store/tests/session_test.rs b/crates/session-store/tests/session_test.rs index d8e27fe0..898bc4bd 100644 --- a/crates/session-store/tests/session_test.rs +++ b/crates/session-store/tests/session_test.rs @@ -16,6 +16,21 @@ use session_store::{FsStore, LogEntry, SegmentStartState, Store, collect_state}; // Helpers // ============================================================================= +fn annotated(items: &[Item]) -> Vec { + items + .iter() + .cloned() + .map(|item| session_store::LoggedHistoryEntry { + item: session_store::LoggedItem::from(item), + metadata: session_store::LoggedSessionHistoryMetadata { + entry_id: session_store::LoggedSessionHistoryEntryId::new(), + origin: session_store::LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + }, + }) + .collect() +} + fn simple_text_events() -> Vec { vec![ Event::text_block_start(0), @@ -144,6 +159,7 @@ async fn run_and_persist( session_id, segment_id, vec![protocol::Segment::text(input)], + annotated(&[Item::user_message(input)]), ) .unwrap(); @@ -154,8 +170,8 @@ async fn run_and_persist( worker.engine = locked.unlock(); let projected = worker.history(); - let new_items = &projected[history_before..]; - session_store::save_delta(store, session_id, segment_id, new_items).unwrap(); + let new_items = annotated(&projected[history_before..]); + session_store::save_delta(store, session_id, segment_id, &new_items).unwrap(); session_store::save_turn_end(store, session_id, segment_id, worker.turn_count()).unwrap(); match &result { @@ -219,7 +235,7 @@ async fn session_run_logs_entries() { SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: &worker.history(), + history: annotated(&worker.history()), }, ) .unwrap(); @@ -237,7 +253,10 @@ async fn session_run_logs_entries() { ); // First entry is SegmentStart - assert!(matches!(&entries[0], LogEntry::SegmentStart { .. })); + assert!(matches!( + &entries[0], + LogEntry::AnnotatedSegmentStart { .. } + )); // Has a RunCompleted with Finished let has_finished = entries.iter().any(|e| { @@ -264,7 +283,7 @@ async fn session_restore_round_trip() { SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: &worker.history(), + history: annotated(&worker.history()), }, ) .unwrap(); @@ -303,7 +322,7 @@ async fn session_run_with_tool_call() { SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: &worker.history(), + history: annotated(&worker.history()), }, ) .unwrap(); @@ -314,12 +333,12 @@ async fn session_run_with_tool_call() { let has_tool_results = entries .iter() - .any(|e| matches!(e, LogEntry::ToolResult { .. })); + .any(|e| matches!(e, LogEntry::AnnotatedToolResult { .. })); assert!(has_tool_results, "should have ToolResult entry"); let has_assistant = entries .iter() - .any(|e| matches!(e, LogEntry::AssistantItem { .. })); + .any(|e| matches!(e, LogEntry::AnnotatedAssistantItem { .. })); assert!(has_assistant, "should have AssistantItem entry"); } @@ -338,7 +357,7 @@ async fn session_resume_after_pause() { SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: &worker.history(), + history: annotated(&worker.history()), }, ) .unwrap(); @@ -377,7 +396,7 @@ async fn session_fork_creates_new_session() { SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: &worker.history(), + history: annotated(&worker.history()), }, ) .unwrap(); @@ -390,7 +409,7 @@ async fn session_fork_creates_new_session() { SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: &worker.history(), + history: annotated(&worker.history()), }, ) .unwrap(); @@ -399,7 +418,10 @@ async fn session_fork_creates_new_session() { // Fork should have a SegmentStart with the current history let fork_entries = store.read_all(fork_sid, fork_segid).unwrap(); assert_eq!(fork_entries.len(), 1); - assert!(matches!(&fork_entries[0], LogEntry::SegmentStart { .. })); + assert!(matches!( + &fork_entries[0], + LogEntry::AnnotatedSegmentStart { .. } + )); let fork_state = collect_state(&fork_entries); assert_eq!(fork_state.session_id, Some(fork_sid)); @@ -418,7 +440,7 @@ async fn session_fork_at_truncates_within_session() { SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: &worker.history(), + history: annotated(&worker.history()), }, ) .unwrap(); @@ -445,6 +467,23 @@ async fn session_fork_at_truncates_within_session() { .expect("source segment has the matching TurnEnd"); let source_state_at_fork = collect_state(&all_entries[..=turn_end_pos]); assert_eq!(fork_state.history.len(), source_state_at_fork.history.len()); + assert_eq!( + fork_state.annotated_history, source_state_at_fork.annotated_history, + "fork_at must preserve every retained history entry identity and provenance", + ); + assert!(fork_state.annotated_history.iter().all(|entry| { + !entry.metadata.entry_id.0.is_empty() + && matches!( + entry.metadata.origin, + session_store::LoggedSessionHistoryOrigin::LegacyUnknown + | session_store::LoggedSessionHistoryOrigin::HumanInput { .. } + | session_store::LoggedSessionHistoryOrigin::WorkerInput { .. } + | session_store::LoggedSessionHistoryOrigin::BackendInstruction { .. } + | session_store::LoggedSessionHistoryOrigin::ModelOutput { .. } + | session_store::LoggedSessionHistoryOrigin::ToolOutput { .. } + | session_store::LoggedSessionHistoryOrigin::DerivedSummary + ) + })); // list_segments should show both source and fork in the same Session. let segs = store.list_segments(sid).unwrap(); @@ -463,7 +502,7 @@ async fn session_config_changed_logged() { SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: &worker.history(), + history: annotated(&worker.history()), }, ) .unwrap(); @@ -496,7 +535,7 @@ async fn session_auto_forks_on_conflict() { SegmentStartState { system_prompt: worker_a.get_system_prompt(), config: worker_a.request_config(), - history: &worker_a.history(), + history: annotated(&worker_a.history()), }, ) .unwrap(); @@ -505,12 +544,14 @@ async fn session_auto_forks_on_conflict() { let mut entries_written: usize = 1; // Simulate another Worker writing to the same segment behind our back. - let extra_entry = LogEntry::UserInput { - ts: 9999, - extensions: vec![], - segments: vec![protocol::Segment::text("Interloper")], - }; - store.append(sid, original_segid, &extra_entry).unwrap(); + session_store::save_user_input( + &store, + sid, + original_segid, + vec![protocol::Segment::text("Interloper")], + annotated(&[Item::user_message("Interloper")]), + ) + .unwrap(); // Now the on-disk count exceeds our tally — ensure_head_or_fork should auto-fork. session_store::ensure_head_or_fork( @@ -522,7 +563,7 @@ async fn session_auto_forks_on_conflict() { SegmentStartState { system_prompt: worker_a.get_system_prompt(), config: worker_a.request_config(), - history: &worker_a.history(), + history: annotated(&worker_a.history()), }, ) .unwrap(); @@ -543,7 +584,7 @@ async fn session_auto_forks_on_conflict() { // The new segment records its lineage forward via forked_from; the // source segment is left immutable (no terminal marker written back). match &fork_entries[0] { - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { forked_from: Some(origin), .. } => { @@ -563,7 +604,7 @@ async fn session_auto_forks_on_conflict() { ); let has_interloper = original_entries .iter() - .any(|e| matches!(e, LogEntry::UserInput { .. })); + .any(|e| matches!(e, LogEntry::AnnotatedUserInput { .. })); assert!(has_interloper); } @@ -581,7 +622,7 @@ async fn nested_past_fork_leaves_ancestors_immutable() { SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: &worker.history(), + history: annotated(&worker.history()), }, ) .unwrap(); @@ -618,7 +659,7 @@ async fn nested_past_fork_leaves_ancestors_immutable() { // fork2's lineage points at fork1, not the root. match &store.read_all(sid, fork2).unwrap()[0] { - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { forked_from: Some(origin), .. } => assert_eq!(origin.segment_id, fork1), diff --git a/crates/tui/src/app.rs b/crates/tui/src/app.rs index ba158e8b..0c620570 100644 --- a/crates/tui/src/app.rs +++ b/crates/tui/src/app.rs @@ -765,7 +765,7 @@ impl App { fn method_for_run(&mut self, segments: Vec) -> Method { // TurnHeader / UserMessage blocks are pushed only after the Worker - // emits `Event::UserMessage` from a committed `LogEntry::UserInput`. + // emits `Event::UserMessage` from a committed `LogEntry::AnnotatedUserInput`. // Locally we only clear the input buffer and forward the method, // while remembering enough local state to undo the visible submit if // the accepted run produced no assistant output and was rolled back. @@ -2937,6 +2937,17 @@ mod composer_history_persistence_tests { mod completion_flow_tests { use super::*; + fn annotated(item: agen::Item) -> session_store::LoggedHistoryEntry { + session_store::LoggedHistoryEntry { + item: session_store::LoggedItem::from(item), + metadata: session_store::LoggedSessionHistoryMetadata { + entry_id: session_store::LoggedSessionHistoryEntryId::new(), + origin: session_store::LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + }, + } + } + #[test] fn typing_at_creates_completion_state_and_emits_query() { let mut app = App::new("test".into()); @@ -3239,7 +3250,7 @@ mod completion_flow_tests { #[test] fn committed_user_message_survives_fresh_segment_rotation() { let mut app = App::new("test".into()); - let start = session_store::LogEntry::SegmentStart { + let start = session_store::LogEntry::AnnotatedSegmentStart { ts: session_store::segment_log::now_millis(), session_id: uuid::Uuid::nil(), system_prompt: None, @@ -3498,14 +3509,14 @@ mod completion_flow_tests { #[test] fn snapshot_excludes_system_prompt_history_from_public_blocks() { let mut app = App::new("test".into()); - let session_start = session_store::LogEntry::SegmentStart { + let session_start = session_store::LogEntry::AnnotatedSegmentStart { ts: 1, session_id: uuid::Uuid::nil(), system_prompt: None, config: Default::default(), - history: vec![session_store::LoggedItem::from( - &agen::Item::system_message("[File: src/main.rs]\nfn main() {}"), - )], + history: vec![annotated(agen::Item::system_message( + "[File: src/main.rs]\nfn main() {}", + ))], forked_from: None, compacted_from: None, }; @@ -3584,7 +3595,7 @@ mod completion_flow_tests { code: ErrorCode::ProviderError, message: "provider unavailable".into(), }); - let segment_start = session_store::LogEntry::SegmentStart { + let segment_start = session_store::LogEntry::AnnotatedSegmentStart { ts: 5, session_id: uuid::Uuid::nil(), system_prompt: None, @@ -4336,36 +4347,33 @@ mod completion_flow_tests { }); let assistant_item_entries = vec![ - serde_json::json!({ - "kind": "assistant_item", - "ts": 1, - "item": { - "kind": "tool_call", - "call_id": "c1", - "name": "TaskCreate", - "arguments": r#"{"subject":"a","description":"A"}"#, - }, - }), - serde_json::json!({ - "kind": "assistant_item", - "ts": 2, - "item": { - "kind": "tool_call", - "call_id": "c2", - "name": "TaskCreate", - "arguments": r#"{"subject":"b","description":"B"}"#, - }, - }), - serde_json::json!({ - "kind": "assistant_item", - "ts": 3, - "item": { - "kind": "tool_call", - "call_id": "u1", - "name": "TaskUpdate", - "arguments": r#"{"taskid":2,"status":"inprogress"}"#, - }, - }), + serde_json::to_value(session_store::LogEntry::AnnotatedAssistantItem { + ts: 1, + entry: annotated(agen::Item::tool_call( + "c1", + "TaskCreate", + r#"{"subject":"a","description":"A"}"#, + )), + }) + .unwrap(), + serde_json::to_value(session_store::LogEntry::AnnotatedAssistantItem { + ts: 2, + entry: annotated(agen::Item::tool_call( + "c2", + "TaskCreate", + r#"{"subject":"b","description":"B"}"#, + )), + }) + .unwrap(), + serde_json::to_value(session_store::LogEntry::AnnotatedAssistantItem { + ts: 3, + entry: annotated(agen::Item::tool_call( + "u1", + "TaskUpdate", + r#"{"taskid":2,"status":"inprogress"}"#, + )), + }) + .unwrap(), ]; app.handle_worker_event(Event::Snapshot { greeting: test_greeting(), diff --git a/crates/tui/src/worker_list.rs b/crates/tui/src/worker_list.rs index fa2d834c..1a6a1be2 100644 --- a/crates/tui/src/worker_list.rs +++ b/crates/tui/src/worker_list.rs @@ -623,6 +623,17 @@ mod tests { const SOURCE: WorkerVisibilitySource = WorkerVisibilitySource::ResumePicker; + fn annotated(item: agen::Item) -> session_store::LoggedHistoryEntry { + session_store::LoggedHistoryEntry { + item: session_store::LoggedItem::from(item), + metadata: session_store::LoggedSessionHistoryMetadata { + entry_id: session_store::LoggedSessionHistoryEntryId::new(), + origin: session_store::LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + }, + } + } + #[test] fn stored_metadata_summary_uses_segment_marker_without_reading_session_log() { let dir = tempdir().unwrap(); @@ -1207,7 +1218,7 @@ mod tests { .append( session_id, segment_id, - &LogEntry::SegmentStart { + &LogEntry::AnnotatedSegmentStart { ts, session_id, system_prompt: None, @@ -1231,9 +1242,10 @@ mod tests { .append( session_id, segment_id, - &LogEntry::UserInput { + &LogEntry::AnnotatedUserInput { ts, segments: vec![protocol::Segment::text(text)], + history: vec![annotated(agen::Item::user_message(text))], extensions: vec![], }, ) diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index d112d7df..e52d0bcf 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -67,8 +67,7 @@ const USER_INPUT_COMMIT_TIMEOUT: Duration = Duration::from_secs(9); fn user_input_has_submission(entry: &LogEntry, submission_id: &str) -> bool { let extensions = match entry { - LogEntry::UserInput { extensions, .. } - | LogEntry::AnnotatedUserInput { extensions, .. } => extensions, + LogEntry::AnnotatedUserInput { extensions, .. } => extensions, _ => return false, }; extensions.iter().any(|extension| { @@ -3244,8 +3243,7 @@ mod tests { assert!(entries.iter().any(|entry| { matches!( entry, - LogEntry::UserInput { segments, .. } - | LogEntry::AnnotatedUserInput { segments, .. } + LogEntry::AnnotatedUserInput { segments, .. } if segments == &vec![Segment::text("start the ticket")] ) })); @@ -3253,8 +3251,7 @@ mod tests { .iter() .find_map(|entry| { let extensions = match entry { - LogEntry::UserInput { extensions, .. } - | LogEntry::AnnotatedUserInput { extensions, .. } => extensions, + LogEntry::AnnotatedUserInput { extensions, .. } => extensions, _ => return None, }; extensions diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 4042ab68..41b33f44 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -631,8 +631,8 @@ fn protocol_command_status(status: WorkdirCommandStatus) -> ProtocolCommandStatu /// /// `Worker::wire_history_persistence` is called separately to wire the /// per-item history commit callback so every assistant / tool item -/// landing in `worker.history` becomes a singular `LogEntry::AssistantItem` -/// / `ToolResult` commit through the sync writer. +/// landing in `worker.history` becomes a singular `LogEntry::AnnotatedAssistantItem` +/// / `AnnotatedToolResult` commit through the sync writer. pub(crate) fn wire_event_bridges_on_engine( worker: &mut Worker, event_tx: &broadcast::Sender, @@ -1317,7 +1317,7 @@ async fn controller_loop( } // Stage the run without a speculative user-message echo. // `Worker::run` validates the input, commits - // `LogEntry::UserInput`, and the session-log sink turns that + // `LogEntry::AnnotatedUserInput`, and the session-log sink turns that // committed entry into the live `Event::UserMessage`. That // keeps every client ordered against `SegmentStart` replay and // makes persisted history the single source of visible user @@ -1343,7 +1343,7 @@ async fn controller_loop( Method::Notify { message, auto_run } => { // Client-side live echo is delivered as `Event::SystemItem` // once the interceptor commits the corresponding - // `LogEntry::SystemItem` entry — drained out of the + // `LogEntry::AnnotatedSystemItem` entry — drained out of the // notify buffer + broadcast through the sink. No // separate echo here. worker.push_notify(message, auto_run); diff --git a/crates/worker/src/feature/builtin/flow_transition.rs b/crates/worker/src/feature/builtin/flow_transition.rs index c307b919..163ac892 100644 --- a/crates/worker/src/feature/builtin/flow_transition.rs +++ b/crates/worker/src/feature/builtin/flow_transition.rs @@ -963,9 +963,12 @@ permission = "read" .append( session_id, segment_id, - &LogEntry::UserInput { + &LogEntry::AnnotatedUserInput { ts: 1, extensions: vec![], + history: vec![crate::session_history::test_logged_history_entry( + agen::Item::user_message("verify current Flow conditions"), + )], segments: vec![Segment::Text { content: "verify current Flow conditions".into(), }], diff --git a/crates/worker/src/hook.rs b/crates/worker/src/hook.rs index c5287b2d..d34ec5cc 100644 --- a/crates/worker/src/hook.rs +++ b/crates/worker/src/hook.rs @@ -161,7 +161,7 @@ impl From for TurnEndAction { /// /// Hook code can use this handle only when the Worker host includes it in an /// event-specific context. The handle queues typed requests; the host drains the -/// queue, commits each entry through `LogEntry::SystemItem`, and only then makes +/// queue, commits each entry through `LogEntry::AnnotatedSystemItem`, and only then makes /// the matching system message visible to the model. It deliberately exposes no /// raw `agen::Item`, history writer, event sender, `Worker`, `Engine`, or /// notification buffer. diff --git a/crates/worker/src/in_flight.rs b/crates/worker/src/in_flight.rs index ecf56f75..095fffb8 100644 --- a/crates/worker/src/in_flight.rs +++ b/crates/worker/src/in_flight.rs @@ -539,9 +539,9 @@ mod tests { text: "done".into(), }], }; - let assistant_entry = LogEntry::AssistantItem { + let assistant_entry = LogEntry::AnnotatedAssistantItem { ts: 1, - item: assistant_item.clone(), + entry: crate::session_history::test_logged_history_entry(assistant_item.clone()), }; let in_flight_guard = in_flight.snapshot_guard(); @@ -593,9 +593,9 @@ mod tests { text: "done".into(), }], }; - let assistant_entry = LogEntry::AssistantItem { + let assistant_entry = LogEntry::AnnotatedAssistantItem { ts: 1, - item: assistant_item.clone(), + entry: crate::session_history::test_logged_history_entry(assistant_item.clone()), }; in_flight.clear_for_committed_item_then(&assistant_item, || { @@ -608,7 +608,7 @@ mod tests { assert!(matches!( entries_snapshot.as_slice(), - [LogEntry::AssistantItem { item, .. }] if item == &assistant_item + [LogEntry::AnnotatedAssistantItem { entry, .. }] if entry.item == assistant_item )); assert!(in_flight_snapshot.is_empty()); } diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index a6283f25..f82444aa 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -60,7 +60,7 @@ pub(crate) struct WorkerInterceptor { pending_notifies: NotifyBuffer, /// Submit-scoped stash of resolver-produced typed system items. /// Drained inside `on_prompt_submit`, committed as - /// `LogEntry::SystemItem` entries through `log_writer`, and + /// `LogEntry::AnnotatedSystemItem` entries through `log_writer`, and /// returned to the worker as `Item::system_message` via /// `PromptAction::ContinueWith`. Populated by `Worker::run` /// immediately before handing off to the worker. @@ -71,7 +71,7 @@ pub(crate) struct WorkerInterceptor { /// Workspace scope associated with Prompt projection provenance. prompt_workspace_id: Option, /// Type-erased commit handle. The interceptor uses it to commit - /// `LogEntry::SystemItem` entries directly (sync) before + /// `LogEntry::AnnotatedSystemItem` entries directly (sync) before /// returning the corresponding `Item::system_message`s up to the /// worker. `None` in tests / `Worker::new` paths where no writer is /// attached. @@ -142,7 +142,7 @@ impl WorkerInterceptor { self } - /// Commit each `SystemItem` as its own `LogEntry::SystemItem` + /// Commit each `SystemItem` as its own `LogEntry::AnnotatedSystemItem` /// entry through the attached writer (no-op when no writer is /// wired). Sync — writes complete before the matching /// `Item::system_message`s reach the worker via @@ -540,7 +540,6 @@ mod tests { entry: session_store::LogEntry, ) -> Result<(), session_store::StoreError> { let item = match entry { - session_store::LogEntry::SystemItem { item, .. } => Some(item), session_store::LogEntry::AnnotatedSystemItem { entry, .. } => Some(entry.item), _ => None, }; diff --git a/crates/worker/src/ipc/notify_buffer.rs b/crates/worker/src/ipc/notify_buffer.rs index 0eb26c66..7b97e4ed 100644 --- a/crates/worker/src/ipc/notify_buffer.rs +++ b/crates/worker/src/ipc/notify_buffer.rs @@ -5,7 +5,7 @@ //! `WorkerInterceptor::pending_history_appends`, which the Engine calls //! at the head of each turn loop iteration. The drain renders each //! pending entry into a typed `SystemItem` (with the `notify_wrapper` -//! prompt applied), commits a `LogEntry::SystemItem` per entry through +//! prompt applied), commits a `LogEntry::AnnotatedSystemItem` per entry through //! the session-log sink, and returns the corresponding //! `Item::system_message`s for the worker to append to its //! persistent history. diff --git a/crates/worker/src/ipc/protocol_session.rs b/crates/worker/src/ipc/protocol_session.rs index 5e16069b..4368fa2b 100644 --- a/crates/worker/src/ipc/protocol_session.rs +++ b/crates/worker/src/ipc/protocol_session.rs @@ -29,18 +29,12 @@ pub fn subscribe_worker_protocol_session(handle: &WorkerHandle) -> WorkerProtoco pub fn live_log_entry_event(entry: LogEntry) -> Option { match entry { - entry @ (LogEntry::SegmentStart { .. } | LogEntry::AnnotatedSegmentStart { .. }) => { + entry @ LogEntry::AnnotatedSegmentStart { .. } => { let session = session_store::public_snapshot::project_current_session_snapshot(&[entry]); Some(Event::SegmentRotated { session }) } - LogEntry::UserInput { segments, .. } | LogEntry::AnnotatedUserInput { segments, .. } => { - Some(Event::UserMessage { segments }) - } - LogEntry::SystemItem { item, .. } => { - let value = serde_json::to_value(&item).expect("SystemItem is Serialize"); - Some(Event::SystemItem { item: value }) - } + LogEntry::AnnotatedUserInput { segments, .. } => Some(Event::UserMessage { segments }), LogEntry::AnnotatedSystemItem { entry, .. } => { let value = serde_json::to_value(&entry.item).expect("SystemItem is Serialize"); Some(Event::SystemItem { item: value }) @@ -89,9 +83,12 @@ mod tests { #[test] fn user_input_log_entry_maps_to_user_message_event() { let segments = vec![protocol::Segment::text("hello from log")]; - let event = live_log_entry_event(LogEntry::UserInput { + let event = live_log_entry_event(LogEntry::AnnotatedUserInput { ts: session_store::segment_log::now_millis(), extensions: vec![], + history: vec![crate::session_history::test_logged_history_entry( + agen::Item::user_message("hello from log"), + )], segments: segments.clone(), }) .expect("UserInput must be live-relevant"); diff --git a/crates/worker/src/ipc/server.rs b/crates/worker/src/ipc/server.rs index 62044ecf..88566597 100644 --- a/crates/worker/src/ipc/server.rs +++ b/crates/worker/src/ipc/server.rs @@ -77,7 +77,7 @@ async fn handle_connection(stream: tokio::net::UnixStream, handle: WorkerHandle) let mut writer = JsonLineWriter::new(writer); // Hold the in-flight stream lock while taking the session-log mirror - // snapshot. `LogEntry::AssistantItem` is mirror-only for live clients, + // snapshot. `LogEntry::AnnotatedAssistantItem` is mirror-only for live clients, // so a finalized assistant block must be observed either as an already // committed entry or as the still-present in-flight block. This lock // order matches `append_entry` (in-flight clear before sink publish) and diff --git a/crates/worker/src/segment_log_sink.rs b/crates/worker/src/segment_log_sink.rs index 77ee51eb..38192af7 100644 --- a/crates/worker/src/segment_log_sink.rs +++ b/crates/worker/src/segment_log_sink.rs @@ -50,7 +50,7 @@ struct SinkInner { /// Broadcast channel for live entry updates. The same `Sender` /// survives session swaps so existing subscribers keep their /// receiver — they observe the swap as a freshly broadcast - /// `LogEntry::SegmentStart` and reset their view accordingly. + /// `LogEntry::AnnotatedSegmentStart` and reset their view accordingly. broadcast_tx: broadcast::Sender, } @@ -89,9 +89,9 @@ impl SegmentLogSink { /// /// Live broadcast fires for committed session-log entries that /// socket clients must see in log order: - /// - `LogEntry::SegmentStart` → `Event::SegmentRotated` on the wire. - /// - `LogEntry::UserInput` → `Event::UserMessage`. - /// - `LogEntry::SystemItem` → `Event::SystemItem`. + /// - `LogEntry::AnnotatedSegmentStart` → `Event::SegmentRotated` on the wire. + /// - `LogEntry::AnnotatedUserInput` → `Event::UserMessage`. + /// - `LogEntry::AnnotatedSystemItem` → `Event::SystemItem`. /// - `LogEntry::Invoke` → `Event::InvokeStart`. /// Everything else (AssistantItem, ToolResult, TurnEnd, /// RunCompleted, RunErrored, PausedTurnAbandoned, LlmUsage, Extension, @@ -120,11 +120,8 @@ impl SegmentLogSink { fn is_live_relevant(entry: &LogEntry) -> bool { matches!( entry, - LogEntry::SegmentStart { .. } - | LogEntry::AnnotatedSegmentStart { .. } - | LogEntry::UserInput { .. } + LogEntry::AnnotatedSegmentStart { .. } | LogEntry::AnnotatedUserInput { .. } - | LogEntry::SystemItem { .. } | LogEntry::AnnotatedSystemItem { .. } | LogEntry::Invoke { .. } ) @@ -132,7 +129,7 @@ impl SegmentLogSink { /// Atomically swap the mirror to `[initial]` and broadcast the new /// session-start entry. Used during compaction / fork: the new - /// `LogEntry::SegmentStart` is the first entry of the replacement + /// `LogEntry::AnnotatedSegmentStart` is the first entry of the replacement /// session, and existing subscribers transition by replaying it /// like any other live entry. /// @@ -234,7 +231,7 @@ mod tests { use session_store::segment_log::now_millis; fn session_start() -> LogEntry { - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { ts: now_millis(), session_id: uuid::Uuid::nil(), system_prompt: None, @@ -253,9 +250,12 @@ mod tests { } fn user_input(text: &str) -> LogEntry { - LogEntry::UserInput { + LogEntry::AnnotatedUserInput { ts: now_millis(), extensions: vec![], + history: vec![crate::session_history::test_logged_history_entry( + agen::Item::user_message(text), + )], segments: vec![protocol::Segment::Text { content: text.to_owned(), }], @@ -270,7 +270,10 @@ mod tests { let (snapshot, mut rx) = sink.subscribe_with_snapshot(); assert_eq!(snapshot.len(), 2); - assert!(matches!(snapshot[0], LogEntry::SegmentStart { .. })); + assert!(matches!( + snapshot[0], + LogEntry::AnnotatedSegmentStart { .. } + )); assert!(matches!( snapshot[1], LogEntry::TurnEnd { turn_count: 1, .. } @@ -279,13 +282,15 @@ mod tests { } fn notification_entry(text: &str) -> LogEntry { - LogEntry::SystemItem { + LogEntry::AnnotatedSystemItem { ts: now_millis(), - item: session_store::SystemItem::Notification { - message: text.to_owned(), - body: format!("[Notification] {text}"), - prompt_provenance: None, - }, + entry: crate::session_history::test_logged_system_entry( + session_store::SystemItem::Notification { + message: text.to_owned(), + body: format!("[Notification] {text}"), + prompt_provenance: None, + }, + ), } } @@ -305,7 +310,7 @@ mod tests { // for Event::UserMessage. sink.publish(user_input("hi from log")); match rx.try_recv() { - Ok(LogEntry::UserInput { segments, .. }) => { + Ok(LogEntry::AnnotatedUserInput { segments, .. }) => { assert_eq!(segments.len(), 1); } other => panic!("expected UserInput, got {other:?}"), @@ -314,7 +319,7 @@ mod tests { // SystemItem is live-relevant. sink.publish(notification_entry("hi")); match rx.try_recv() { - Ok(LogEntry::SystemItem { .. }) => {} + Ok(LogEntry::AnnotatedSystemItem { .. }) => {} other => panic!("expected SystemItem, got {other:?}"), } @@ -332,7 +337,7 @@ mod tests { assert_eq!(snapshot.len(), 1); match rx.try_recv() { - Ok(LogEntry::SystemItem { .. }) => {} + Ok(LogEntry::AnnotatedSystemItem { .. }) => {} other => panic!("unexpected: {other:?}"), } assert!(rx.try_recv().is_err()); @@ -348,13 +353,16 @@ mod tests { sink.reset_with_initial(session_start()); match rx.try_recv() { - Ok(LogEntry::SegmentStart { .. }) => {} + Ok(LogEntry::AnnotatedSegmentStart { .. }) => {} other => panic!("expected SegmentStart broadcast, got {other:?}"), } let (post_snapshot, _) = sink.subscribe_with_snapshot(); assert_eq!(post_snapshot.len(), 1); - assert!(matches!(post_snapshot[0], LogEntry::SegmentStart { .. })); + assert!(matches!( + post_snapshot[0], + LogEntry::AnnotatedSegmentStart { .. } + )); } #[test] diff --git a/crates/worker/src/session_history.rs b/crates/worker/src/session_history.rs index 3772d554..33ac6eed 100644 --- a/crates/worker/src/session_history.rs +++ b/crates/worker/src/session_history.rs @@ -5,7 +5,6 @@ //! retained only as explicit `LegacyUnknown` entries. use agen::{HistoryEntry, Item}; -use protocol::Segment; use session_store::{ LogEntry, LoggedHistoryDerivation, LoggedHistoryEntry, LoggedSessionHistoryEntryId, LoggedSessionHistoryMetadata, LoggedSessionHistoryOrigin, LoggedWorkerSubject, SegmentId, @@ -18,6 +17,32 @@ pub type WorkerHistoryProvenance = LoggedSessionHistoryOrigin; pub type SessionHistoryDerivation = LoggedHistoryDerivation; pub type WorkerSubjectSnapshot = LoggedWorkerSubject; +#[cfg(test)] +pub(crate) fn test_logged_history_entry(item: impl Into) -> LoggedHistoryEntry { + LoggedHistoryEntry { + item: session_store::LoggedItem::from(item.into()), + metadata: LoggedSessionHistoryMetadata { + entry_id: LoggedSessionHistoryEntryId::new(), + origin: LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + }, + } +} + +#[cfg(test)] +pub(crate) fn test_logged_system_entry( + item: session_store::SystemItem, +) -> session_store::LoggedSystemHistoryEntry { + session_store::LoggedSystemHistoryEntry { + item, + metadata: LoggedSessionHistoryMetadata { + entry_id: LoggedSessionHistoryEntryId::new(), + origin: LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + }, + } +} + pub(crate) fn worker_subject(session_id: SessionId) -> WorkerSubjectSnapshot { WorkerSubjectSnapshot { workspace_id: None, @@ -53,10 +78,6 @@ pub(crate) fn to_logged_history_entry( } } -fn legacy_entry(item: Item) -> HistoryEntry { - HistoryEntry::new(item, SessionHistoryMetadata::legacy_unknown()) -} - fn from_logged(entry: &LoggedHistoryEntry) -> HistoryEntry { HistoryEntry::new(Item::from(entry.item.clone()), entry.metadata.clone()) } @@ -74,32 +95,15 @@ pub(crate) fn restore_history_entries( LogEntry::AnnotatedSegmentStart { history: seed, .. } => { history = seed.iter().map(from_logged).collect(); } - LogEntry::SegmentStart { history: seed, .. } => { - history = seed - .iter() - .cloned() - .map(Item::from) - .map(legacy_entry) - .collect(); - } LogEntry::AnnotatedUserInput { history: input, .. } => { history.extend(input.iter().map(from_logged)) } - LogEntry::UserInput { segments, .. } => history.push(legacy_entry(Item::user_message( - Segment::flatten_to_text(segments), - ))), LogEntry::AnnotatedAssistantItem { entry, .. } | LogEntry::AnnotatedToolResult { entry, .. } => history.push(from_logged(entry)), - LogEntry::AssistantItem { item, .. } | LogEntry::ToolResult { item, .. } => { - history.push(legacy_entry(Item::from(item.clone()))); - } LogEntry::AnnotatedSystemItem { entry, .. } => history.push(HistoryEntry::new( entry.item.to_history_item(), entry.metadata.clone(), )), - LogEntry::SystemItem { item, .. } => { - history.push(legacy_entry(item.to_history_item())); - } _ => {} } } @@ -110,23 +114,9 @@ pub(crate) fn restore_history_entries( mod tests { use super::*; use agen::llm_client::RequestConfig; + use protocol::Segment; use session_store::LogEntry; - #[test] - fn legacy_user_role_is_not_inferred_as_human_authority() { - let entries = vec![LogEntry::UserInput { - ts: 1, - segments: vec![Segment::text("legacy")], - extensions: Vec::new(), - }]; - let restored = - restore_history_entries(SessionId::now_v7(), SegmentId::now_v7(), &entries).unwrap(); - assert!(matches!( - restored[0].annotation.origin, - WorkerHistoryProvenance::LegacyUnknown - )); - } - #[test] fn typed_flow_and_unknown_caller_input_round_trip_without_role_inference() { let session_id = SessionId::now_v7(); diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index 8e0ea0f1..2c769f6e 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -112,8 +112,12 @@ impl InternalSpawnedWorkerRecord { fn stop_summary(&self) -> SubWorkerStopSummary { let mut counts = BTreeMap::::new(); for entry in self.session.entries() { - if let session_store::LogEntry::AssistantItem { - item: LoggedItem::ToolCall { name, .. }, + if let session_store::LogEntry::AnnotatedAssistantItem { + entry: + session_store::LoggedHistoryEntry { + item: LoggedItem::ToolCall { name, .. }, + .. + }, .. } = entry { @@ -1024,11 +1028,16 @@ mod tests { && worker.parent_session_id.as_deref() == Some("parent-session") && matches!(*event, Event::TextDone { ref text } if text == "answer") )); - record.session.publish_test_entry(LogEntry::UserInput { - ts: 1, - segments: vec![protocol::Segment::text("question")], - extensions: Vec::new(), - }); + record + .session + .publish_test_entry(LogEntry::AnnotatedUserInput { + ts: 1, + segments: vec![protocol::Segment::text("question")], + history: vec![crate::session_history::test_logged_history_entry( + agen::Item::user_message("question"), + )], + extensions: Vec::new(), + }); let committed = tokio::time::timeout(Duration::from_secs(1), parent_rx.recv()) .await .unwrap() @@ -1159,14 +1168,18 @@ mod tests { let (mut record, _events) = record("child", InternalWorkerVisibility::ParentClient).await; record.change_tracker = Some(tracker); for (index, name) in ["Read", "Read", "Grep"].into_iter().enumerate() { - record.session.publish_test_entry(LogEntry::AssistantItem { - ts: index as u64, - item: LoggedItem::ToolCall { - call_id: format!("call-{index}"), - name: name.to_string(), - arguments: "{}".to_string(), - }, - }); + record + .session + .publish_test_entry(LogEntry::AnnotatedAssistantItem { + ts: index as u64, + entry: crate::session_history::test_logged_history_entry( + LoggedItem::ToolCall { + call_id: format!("call-{index}"), + name: name.to_string(), + arguments: "{}".to_string(), + }, + ), + }); } registry.start_protocol_forwarding(record.clone()); install_record(®istry, record); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 28f69623..80dd85b2 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -900,7 +900,6 @@ where self.state.increment_entries(); if let Some(in_flight) = &self.in_flight { let committed_item = match &entry { - LogEntry::AssistantItem { item, .. } => Some(item.clone()), LogEntry::AnnotatedAssistantItem { entry, .. } => Some(entry.item.clone()), _ => None, }; @@ -1207,11 +1206,11 @@ pub struct Worker { memory_task: Option>, /// Typed user submissions in submit order. K-th entry corresponds to /// the K-th `Item::user_message` in `worker.history()` (modulo seed - /// history loaded via `SegmentStart.history`, whose original segments + /// history loaded via `AnnotatedSegmentStart.history`, whose original segments /// are not preserved). Populated from log on `restore_from_manifest`, /// appended after `save_user_input` on each `run`. Pre-`Event::Snapshot` /// this fed `WorkerSharedState.user_segments`; the new wire format - /// carries typed atoms via `LogEntry::UserInput { segments }` so + /// carries typed atoms via `LogEntry::AnnotatedUserInput { segments }` so /// this remains purely an in-memory tracker for compact alignment. user_segments: Vec>, /// Worker-side session-log mirror + broadcast sink. Populated alongside @@ -1221,7 +1220,8 @@ pub struct Worker { sink: SegmentLogSink, /// `true` once `wire_history_persistence` has installed the /// `Engine::on_history_append` callback that commits each appended - /// item as a singular `LogEntry::AssistantItem` / `ToolResult` + /// item as a singular `LogEntry::AnnotatedAssistantItem` / + /// `AnnotatedToolResult` /// directly through the writer. Tests that drive `Worker::new` without /// going through the controller leave this `false`; `persist_turn` /// then walks the post-`history_before` slice inline so entries @@ -1345,15 +1345,16 @@ impl Worker { } /// Wire `Engine::on_history_append` to commit each appended item - /// directly as a singular `LogEntry::AssistantItem` / `ToolResult` + /// directly as a singular `LogEntry::AnnotatedAssistantItem` / + /// `AnnotatedToolResult` /// through the writer. The controller calls this once per spawned /// Worker after the worker is built; tests that drive `Worker::new` may /// opt in to the same wiring or leave it off (in which case /// `persist_turn`'s inline fallback writes entries at turn end). /// /// `user_message` items are skipped because they are committed - /// up-front via `commit_entry(LogEntry::UserInput { segments })`. - /// `role:system` items are committed as typed `LogEntry::SystemItem` + /// up-front via `commit_entry(LogEntry::AnnotatedUserInput { segments })`. + /// `role:system` items are committed as typed `LogEntry::AnnotatedSystemItem` /// entries by their producers (for example `WorkerInterceptor` and /// interrupted-turn prep) before they reach the worker's history, so this /// callback would otherwise double-write them. @@ -1941,8 +1942,7 @@ impl Worker { } let input = match entries.get(target.user_input_entry_index) { - Some(LogEntry::UserInput { segments, .. }) - | Some(LogEntry::AnnotatedUserInput { segments, .. }) => segments.clone(), + Some(LogEntry::AnnotatedUserInput { segments, .. }) => segments.clone(), _ => { return Err(RewindError::Invalid( "rewind target is no longer a user message".into(), @@ -2081,8 +2081,8 @@ impl Worker { /// Cheap clone via `Option`. /// Snapshot of the typed user segments tracked alongside worker /// history. The K-th entry corresponds to the K-th `Item::user_message` - /// derived from `LogEntry::UserInput` entries (post-compaction); seed - /// history loaded via `SegmentStart.history` does not contribute, + /// derived from `LogEntry::AnnotatedUserInput` entries (post-compaction); seed + /// history loaded via `AnnotatedSegmentStart.history` does not contribute, /// which is acceptable because the original segments are unrecoverable. pub fn user_segments(&self) -> &[Vec] { &self.user_segments @@ -3533,12 +3533,13 @@ impl Worker { // slice from `history_before` inline so the test's // `restore`-style assertions still see entries on disk. if !self.history_persistence_wired { - let new_items: Vec = self.session.history().entries()[history_before..] + let new_entries: Vec<_> = self.session.history().entries()[history_before..] .iter() - .map(|entry| entry.item.clone()) + .map(to_logged_history_entry) .collect(); let ts = segment_log::now_millis(); - for item in &new_items { + for history_entry in new_entries { + let item = Item::from(history_entry.item.clone()); if item.is_user_message() { continue; } @@ -3551,7 +3552,7 @@ impl Worker { ) { continue; } - let entry = session_store::classify_history_item(item, ts); + let entry = session_store::classify_logged_history_entry(history_entry, ts); self.commit_entry(entry)?; } } @@ -6257,8 +6258,7 @@ fn build_rewind_targets(segment_id: uuid::Uuid, entries: &[LogEntry]) -> Vec (segments, ts), + LogEntry::AnnotatedUserInput { segments, ts, .. } => (segments, ts), _ => continue, }; turn_index += 1; @@ -6300,8 +6300,7 @@ fn rewind_truncate_entries(entries: &[LogEntry], user_input_entry_index: usize) fn suffix_has_tool_side_effects(entries: &[LogEntry]) -> bool { entries.iter().any(|entry| match entry { - LogEntry::ToolResult { .. } | LogEntry::AnnotatedToolResult { .. } => true, - LogEntry::AssistantItem { item, .. } => logged_item_is_tool_call(item), + LogEntry::AnnotatedToolResult { .. } => true, LogEntry::AnnotatedAssistantItem { entry, .. } => logged_item_is_tool_call(&entry.item), _ => false, }) @@ -7636,7 +7635,7 @@ mod build_summary_prompt_tests { ); assert!(checkpoint.is_none()); - let mut replacement_entries = vec![LogEntry::SegmentStart { + let mut replacement_entries = vec![LogEntry::AnnotatedSegmentStart { ts: segment_log::now_millis(), session_id: uuid::Uuid::nil(), system_prompt: None, @@ -7964,9 +7963,12 @@ mod build_summary_prompt_tests { ); append_test_entry( worker, - LogEntry::UserInput { + LogEntry::AnnotatedUserInput { ts: ts + 1, extensions: vec![], + history: vec![crate::session_history::test_logged_history_entry( + Item::user_message(text), + )], segments: vec![text_segment(text)], }, ); @@ -7986,16 +7988,18 @@ mod build_summary_prompt_tests { append_user_turn(&worker, 20, "second message"); append_test_entry( &worker, - LogEntry::ToolResult { + LogEntry::AnnotatedToolResult { ts: 30, - item: session_store::LoggedItem::ToolResult { - call_id: "call-1".into(), - summary: "wrote a file".into(), - content: None, - attachments: Vec::new(), - disposition: Default::default(), - is_error: false, - }, + entry: crate::session_history::test_logged_history_entry( + session_store::LoggedItem::ToolResult { + call_id: "call-1".into(), + summary: "wrote a file".into(), + content: None, + attachments: Vec::new(), + disposition: Default::default(), + is_error: false, + }, + ), }, ); @@ -8029,16 +8033,18 @@ mod build_summary_prompt_tests { append_user_turn(&worker, 20, "second message"); append_test_entry( &worker, - LogEntry::ToolResult { + LogEntry::AnnotatedToolResult { ts: 30, - item: session_store::LoggedItem::ToolResult { - call_id: "call-1".into(), - summary: "wrote a file".into(), - content: None, - attachments: Vec::new(), - disposition: Default::default(), - is_error: false, - }, + entry: crate::session_history::test_logged_history_entry( + session_store::LoggedItem::ToolResult { + call_id: "call-1".into(), + summary: "wrote a file".into(), + content: None, + attachments: Vec::new(), + disposition: Default::default(), + is_error: false, + }, + ), }, ); let (head_entries, targets) = worker.list_rewind_targets().unwrap(); @@ -8456,9 +8462,9 @@ mod build_summary_prompt_tests { worker.wire_history_persistence(); let dangling_call = Item::tool_call("call-1", "SideEffect", "{}"); worker - .commit_entry(LogEntry::AssistantItem { + .commit_entry(LogEntry::AnnotatedAssistantItem { ts: segment_log::now_millis(), - item: dangling_call.clone().into(), + entry: crate::session_history::test_logged_history_entry(dangling_call.clone()), }) .unwrap(); worker.set_history_for_test(vec![dangling_call]); @@ -8863,9 +8869,12 @@ mod build_summary_prompt_tests { ); worker.set_history_for_test(vec![evidence.clone()]); worker - .commit_entry(LogEntry::UserInput { + .commit_entry(LogEntry::AnnotatedUserInput { ts: segment_log::now_millis(), extensions: vec![], + history: vec![crate::session_history::test_logged_history_entry( + evidence.clone(), + )], segments: vec![text_segment( "The cancellation regression must leave this evidence available for retry.", )], diff --git a/crates/worker/tests/compact_events_test.rs b/crates/worker/tests/compact_events_test.rs index 67e81367..e67519fe 100644 --- a/crates/worker/tests/compact_events_test.rs +++ b/crates/worker/tests/compact_events_test.rs @@ -25,6 +25,17 @@ use worker::{Worker, WorkerController}; type TestStore = CombinedStore; +fn annotated(item: Item) -> session_store::LoggedHistoryEntry { + session_store::LoggedHistoryEntry { + item: session_store::LoggedItem::from(item), + metadata: session_store::LoggedSessionHistoryMetadata { + entry_id: session_store::LoggedSessionHistoryEntryId::new(), + origin: session_store::LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + }, + } +} + #[derive(Clone)] struct MockClient { responses: Arc>>, @@ -210,7 +221,6 @@ fn system_texts_in_sink_session_start( .into_iter() .map(|entry| entry.item) .collect::>(), - session_store::LogEntry::SegmentStart { history, .. } => history, _ => continue, }; return history @@ -310,17 +320,14 @@ permission = "write" // Simulate a foreign writer appending to the same segment. This bumps // the on-disk entry count past the Worker's own append tally without // updating the Worker's `entries_written`. - store - .append( - session_id, - source_segment_id, - &LogEntry::UserInput { - ts: 9999, - segments: vec![protocol::Segment::text("interloper")], - extensions: vec![], - }, - ) - .unwrap(); + session_store::save_user_input( + &store, + session_id, + source_segment_id, + vec![protocol::Segment::text("interloper")], + vec![annotated(Item::user_message("interloper"))], + ) + .unwrap(); // Next run triggers ensure_segment_head, which sees the drift. worker.run_text("second").await.unwrap(); @@ -348,11 +355,6 @@ permission = "write" session_id: seg_session, forked_from: Some(origin), .. - } - | LogEntry::SegmentStart { - session_id: seg_session, - forked_from: Some(origin), - .. } => { assert_eq!(*seg_session, session_id); assert_eq!(origin.segment_id, source_segment_id); @@ -366,7 +368,7 @@ permission = "write" assert_eq!(source_after.len(), source_len_before + 1); assert!(matches!( source_after.last(), - Some(LogEntry::UserInput { .. }) + Some(LogEntry::AnnotatedUserInput { .. }) )); } diff --git a/crates/worker/tests/controller_test.rs b/crates/worker/tests/controller_test.rs index 078fdc97..92191f3e 100644 --- a/crates/worker/tests/controller_test.rs +++ b/crates/worker/tests/controller_test.rs @@ -35,29 +35,16 @@ fn history_from_sink(handle: &WorkerHandle) -> Vec { LogEntry::AnnotatedSegmentStart { history, .. } => { items.extend(history.into_iter().map(|entry| Item::from(entry.item))); } - LogEntry::SegmentStart { history, .. } => { - items.extend(history.into_iter().map(Item::from)); - } LogEntry::AnnotatedUserInput { history, .. } => { items.extend(history.into_iter().map(|entry| Item::from(entry.item))); } - LogEntry::UserInput { segments, .. } => { - let text = protocol::Segment::flatten_to_text(&segments); - items.push(Item::user_message(text)); - } LogEntry::AnnotatedAssistantItem { entry, .. } | LogEntry::AnnotatedToolResult { entry, .. } => { items.push(Item::from(entry.item)); } - LogEntry::AssistantItem { item, .. } | LogEntry::ToolResult { item, .. } => { - items.push(Item::from(item)); - } LogEntry::AnnotatedSystemItem { entry, .. } => { items.push(entry.item.to_history_item()); } - LogEntry::SystemItem { item, .. } => { - items.push(item.to_history_item()); - } _ => {} } } @@ -67,7 +54,6 @@ fn history_from_sink(handle: &WorkerHandle) -> Vec { fn system_item(entry: &LogEntry) -> Option<&session_store::SystemItem> { match entry { LogEntry::AnnotatedSystemItem { entry, .. } => Some(&entry.item), - LogEntry::SystemItem { item, .. } => Some(item), _ => None, } } @@ -1095,7 +1081,7 @@ async fn run_with_paste_segment_inlines_content_and_emits_typed_user_message() { // Mixed input: plain text + a paste chip + trailing text. Worker must // flatten this into one user-message string (paste content inlined, // no `[Clipboard ...]` label leaking to the LLM); the committed - // `LogEntry::UserInput` must carry the typed segments unchanged so + // `LogEntry::AnnotatedUserInput` must carry the typed segments unchanged so // socket clients can derive `Event::UserMessage` and re-render the chip. let segments = vec![ protocol::Segment::text("see "), @@ -1130,7 +1116,7 @@ async fn run_with_paste_segment_inlines_content_and_emits_typed_user_message() { _ => {} }, entry = entry_rx.recv() => match entry { - Ok(session_store::LogEntry::UserInput { segments, .. } | session_store::LogEntry::AnnotatedUserInput { segments, .. }) => { + Ok(session_store::LogEntry::AnnotatedUserInput { segments, .. }) => { user_input_segments = Some(segments); if saw_turn_end { break; From adb684a6bf80305954819f0240d2dda75cb603ed Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 30 Aug 2026 12:33:15 +0900 Subject: [PATCH 4/4] fix: preserve typed user input snapshots --- crates/session-store/src/public_snapshot.rs | 103 ++++++++++++++++++-- 1 file changed, 97 insertions(+), 6 deletions(-) diff --git a/crates/session-store/src/public_snapshot.rs b/crates/session-store/src/public_snapshot.rs index 3109bbe3..3fdb32e1 100644 --- a/crates/session-store/src/public_snapshot.rs +++ b/crates/session-store/src/public_snapshot.rs @@ -94,12 +94,14 @@ fn extend_history( let mut attached_segments = false; for entry in history { let data = if !attached_segments - && matches!( - entry.metadata.origin, - LoggedSessionHistoryOrigin::HumanInput { .. } - ) && input_segments.is_some() - { + && matches!( + &entry.item, + LoggedItem::Message { + role: LoggedRole::User, + .. + } + ) { attached_segments = true; SessionSnapshotEntryData::UserInput { segments: input_segments.cloned().unwrap_or_default(), @@ -293,7 +295,10 @@ mod tests { use agen::llm_client::RequestConfig; use super::*; - use crate::{LoggedSessionHistoryEntryId, LoggedSessionHistoryMetadata, LoggedWorkerSubject}; + use crate::{ + LoggedHistoryDerivation, LoggedSessionHistoryEntryId, LoggedSessionHistoryMetadata, + LoggedWorkerSubject, + }; #[test] fn current_projection_is_stable_and_hides_reasoning_and_system_prompts() { @@ -352,6 +357,92 @@ mod tests { assert!(json.contains("visible")); } + #[test] + fn annotated_user_input_attaches_segments_to_first_user_role_entry_for_any_origin() { + let session_id = crate::new_session_id(); + let segments = vec![Segment::Text { + content: "normal submit".into(), + }]; + + for origin in [ + LoggedSessionHistoryOrigin::LegacyUnknown, + LoggedSessionHistoryOrigin::FlowInstruction { + selector: "builtin:coder-review".into(), + definition_id: "flow-definition".into(), + definition_revision: 7, + instance_id: "flow-instance".into(), + state_id: "implement".into(), + }, + ] { + let user_entry_id = LoggedSessionHistoryEntryId::new(); + let source_entry_id = LoggedSessionHistoryEntryId::new(); + let log = vec![ + LogEntry::AnnotatedSegmentStart { + ts: 1, + session_id, + system_prompt: None, + config: RequestConfig::default(), + history: Vec::new(), + forked_from: None, + compacted_from: None, + }, + LogEntry::AnnotatedUserInput { + ts: 2, + segments: segments.clone(), + history: vec![ + LoggedHistoryEntry { + item: LoggedItem::Message { + role: LoggedRole::System, + content: vec![LoggedContentPart::Text { + text: "flow instruction".into(), + }], + }, + metadata: LoggedSessionHistoryMetadata { + entry_id: LoggedSessionHistoryEntryId::new(), + origin: LoggedSessionHistoryOrigin::FlowInstruction { + selector: "builtin:coder-review".into(), + definition_id: "flow-definition".into(), + definition_revision: 7, + instance_id: "flow-instance".into(), + state_id: "implement".into(), + }, + derivation: None, + }, + }, + LoggedHistoryEntry { + item: LoggedItem::Message { + role: LoggedRole::User, + content: vec![LoggedContentPart::Text { + text: "normal submit".into(), + }], + }, + metadata: LoggedSessionHistoryMetadata { + entry_id: user_entry_id.clone(), + origin: origin.clone(), + derivation: Some(LoggedHistoryDerivation { + sources: vec![source_entry_id.clone()], + }), + }, + }, + ], + extensions: Vec::new(), + }, + ]; + + let snapshot = project_current_session_snapshot(&log); + assert_eq!(snapshot.entries.len(), 1); + assert_eq!(snapshot.entries[0].entry_id, user_entry_id.0); + assert_eq!(snapshot.entries[0].provenance, provenance(&origin)); + assert_eq!(snapshot.entries[0].derived_from, vec![source_entry_id.0]); + assert_eq!( + snapshot.entries[0].data, + SessionSnapshotEntryData::UserInput { + segments: segments.clone(), + } + ); + } + } + #[test] fn annotated_projection_preserves_identity_and_provenance() { let session_id = crate::new_session_id();