diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index d0d6e5a4..8ca42409 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -1416,6 +1416,28 @@ pub enum CommandEvent { }, } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +pub struct InFlightCompaction { + pub schema_version: u32, + pub compaction_id: String, + pub revision: u64, + pub internal_worker: Option, + pub started_at_ms: u64, +} + +impl InFlightCompaction { + pub fn from_running(lifecycle: &CompactionLifecycle) -> Option { + (lifecycle.state == CompactionLifecycleState::Running).then(|| Self { + schema_version: lifecycle.schema_version, + compaction_id: lifecycle.compaction_id.clone(), + revision: lifecycle.revision, + internal_worker: lifecycle.internal_worker.clone(), + started_at_ms: lifecycle.started_at_ms, + }) + } +} + /// Unfinished model output and active command state included in /// `Event::Snapshot` for clients that attach while work is still streaming. /// @@ -1430,11 +1452,17 @@ pub struct InFlightSnapshot { pub blocks: Vec, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub commands: Vec, + /// The currently running compaction, if any. + /// + /// This is lifecycle progress only. Candidate history and the staged + /// Segment remain private until the Segment is activated atomically. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub compaction: Option, } impl InFlightSnapshot { pub fn is_empty(&self) -> bool { - self.blocks.is_empty() && self.commands.is_empty() + self.blocks.is_empty() && self.commands.is_empty() && self.compaction.is_none() } } @@ -2280,6 +2308,13 @@ mod tests { stderr: CommandStreamSlice::default(), exit_code: None, }], + compaction: Some(InFlightCompaction { + schema_version: 3, + compaction_id: "compaction-1".into(), + revision: 1, + internal_worker: None, + started_at_ms: 99, + }), }, internal_workers: Vec::new(), }; @@ -2291,9 +2326,26 @@ mod tests { parsed["data"]["in_flight"]["blocks"][2]["state"], "streaming_args" ); + assert_eq!( + parsed["data"]["in_flight"]["compaction"]["compaction_id"], + "compaction-1" + ); + assert!( + parsed["data"]["in_flight"]["compaction"] + .as_object() + .is_some_and(|value| { + !value.contains_key("state") + && !value.contains_key("summary") + && !value.contains_key("new_segment_id") + }), + "in-flight compaction progress must not expose terminal or staged state" + ); match serde_json::from_str::(&json).unwrap() { - Event::Snapshot { in_flight, .. } => assert_eq!(in_flight.blocks.len(), 3), + Event::Snapshot { in_flight, .. } => { + assert_eq!(in_flight.blocks.len(), 3); + assert_eq!(in_flight.compaction.unwrap().compaction_id, "compaction-1"); + } other => panic!("expected Snapshot, got {other:?}"), } } diff --git a/crates/protocol/src/typescript.rs b/crates/protocol/src/typescript.rs index f3f147a4..b5fec423 100644 --- a/crates/protocol/src/typescript.rs +++ b/crates/protocol/src/typescript.rs @@ -5,15 +5,15 @@ use ts_rs::{Config, TS}; use crate::{ Alert, AlertLevel, AlertSource, CommandEvent, CommandSnapshot, CommandStatus, CommandStream, CommandStreamSlice, CompactionLifecycle, CompactionLifecycleState, CompletionEntry, - CompletionKind, ErrorCode, Event, Greeting, InFlightBlock, InFlightSnapshot, - InFlightToolCallState, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnapshot, - InvokeKind, MemoryWorkerEvent, Method, PasteArtifactAvailability, PasteArtifactMediaType, - PasteArtifactRef, PendingSubmissionSummary, PendingSubmissionsSnapshot, Permission, - RewindSummary, RewindTarget, RewindTargetId, RunResult, ScopeRule, Segment, SessionContentPart, - SessionEntryProvenance, SessionMessageRole, SessionSnapshot, SessionSnapshotEntry, - SessionSnapshotEntryData, SessionToolAttachment, SubmissionDisposition, SymlinkPolicy, - ToolResultDisposition, TurnResult, UploadedFileAvailability, UploadedFileRef, WorkerBusyState, - WorkerCommandAcknowledgement, WorkerCommandDisposition, WorkerCommandEnvelope, + CompletionKind, ErrorCode, Event, Greeting, InFlightBlock, InFlightCompaction, + InFlightSnapshot, InFlightToolCallState, InternalWorkerKind, InternalWorkerRef, + InternalWorkerSnapshot, InvokeKind, MemoryWorkerEvent, Method, PasteArtifactAvailability, + PasteArtifactMediaType, PasteArtifactRef, PendingSubmissionSummary, PendingSubmissionsSnapshot, + Permission, RewindSummary, RewindTarget, RewindTargetId, RunResult, ScopeRule, Segment, + SessionContentPart, SessionEntryProvenance, SessionMessageRole, SessionSnapshot, + SessionSnapshotEntry, SessionSnapshotEntryData, SessionToolAttachment, SubmissionDisposition, + SymlinkPolicy, ToolResultDisposition, TurnResult, UploadedFileAvailability, UploadedFileRef, + WorkerBusyState, WorkerCommandAcknowledgement, WorkerCommandDisposition, WorkerCommandEnvelope, WorkerCommandKind, WorkerEvent, WorkerMaintenanceState, WorkerRunState, WorkerState, WorkerStateSnapshot, WorkerStatus, subscription::{ @@ -81,6 +81,7 @@ 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); diff --git a/crates/tui/src/app.rs b/crates/tui/src/app.rs index 9858dd39..b0483bf6 100644 --- a/crates/tui/src/app.rs +++ b/crates/tui/src/app.rs @@ -277,6 +277,8 @@ pub struct App { /// Turn/protocol errors retained when a real `SegmentStart` replaces the /// replayable conversation rows during segment rotation. run_error_messages: Vec, + /// Current compaction identity/revision used to fence snapshot/live updates. + active_compaction: Option<(String, u64)>, /// Presentation-only Internal Worker projections keyed by session identity. /// They are rendered in separate selectable views and never mixed into `blocks`. pub internal_workers: Vec, @@ -364,6 +366,7 @@ impl App { quit_confirm: None, shutdown_confirm: None, blocks: Vec::new(), + active_compaction: None, run_error_messages: Vec::new(), internal_workers: Vec::new(), selected_internal_worker_session_id: None, @@ -1397,14 +1400,33 @@ impl App { self.reset_run_state(); } } - Event::CompactStart { .. } => { - if self.last_streaming_compact_mut().is_none() { - self.blocks.push(Block::Compact(CompactEvent::Streaming { - started_at: Instant::now(), - })); + Event::CompactStart { lifecycle } => { + let should_apply = match &self.active_compaction { + None => true, + Some((id, revision)) => { + id == &lifecycle.compaction_id && lifecycle.revision > *revision + } + }; + if should_apply { + self.active_compaction = Some((lifecycle.compaction_id, lifecycle.revision)); + if self.last_streaming_compact_mut().is_none() { + self.blocks.push(Block::Compact(CompactEvent::Streaming { + started_at: Instant::now(), + })); + } } } Event::CompactDone { lifecycle } => { + let should_apply = match &self.active_compaction { + None => true, + Some((id, revision)) => { + id == &lifecycle.compaction_id && lifecycle.revision > *revision + } + }; + if !should_apply { + return None; + } + self.active_compaction = None; self.session_context_tokens = 0; let new_segment_id = lifecycle .new_segment_id @@ -1430,6 +1452,16 @@ impl App { } } Event::CompactFailed { lifecycle } => { + let should_apply = match &self.active_compaction { + None => true, + Some((id, revision)) => { + id == &lifecycle.compaction_id && lifecycle.revision > *revision + } + }; + if !should_apply { + return None; + } + self.active_compaction = None; let error = lifecycle .error .unwrap_or_else(|| "compaction failed".to_string()); @@ -1614,6 +1646,7 @@ impl App { } fn apply_in_flight_snapshot(&mut self, snapshot: InFlightSnapshot) { + let compaction = snapshot.compaction; for block in snapshot.blocks { match block { InFlightBlock::Text { text, finished } => { @@ -1655,6 +1688,14 @@ impl App { } } } + self.active_compaction = compaction + .as_ref() + .map(|lifecycle| (lifecycle.compaction_id.clone(), lifecycle.revision)); + if compaction.is_some() && self.last_streaming_compact_mut().is_none() { + self.blocks.push(Block::Compact(CompactEvent::Streaming { + started_at: Instant::now(), + })); + } } fn append_assistant_text(&mut self, text: &str) { @@ -3773,6 +3814,7 @@ mod completion_flow_tests { }, ], commands: Vec::new(), + compaction: None, }, internal_workers: Vec::new(), }); @@ -4222,6 +4264,7 @@ mod completion_flow_tests { lifecycle: test_compaction_lifecycle(protocol::CompactionLifecycleState::Running), }); let mut lifecycle = test_compaction_lifecycle(protocol::CompactionLifecycleState::Done); + lifecycle.revision = 2; lifecycle.new_segment_id = Some(id.to_string()); app.handle_worker_event(Event::CompactDone { lifecycle }); @@ -4243,6 +4286,7 @@ mod completion_flow_tests { lifecycle: test_compaction_lifecycle(protocol::CompactionLifecycleState::Running), }); let mut lifecycle = test_compaction_lifecycle(protocol::CompactionLifecycleState::Failed); + lifecycle.revision = 2; lifecycle.error = Some("provider 429".into()); app.handle_worker_event(Event::CompactFailed { lifecycle }); @@ -4256,6 +4300,31 @@ mod completion_flow_tests { )); } + #[test] + fn snapshot_restores_running_compaction_and_fences_unrelated_terminal() { + let mut app = App::new("test".into()); + let lifecycle = test_compaction_lifecycle(protocol::CompactionLifecycleState::Running); + app.apply_in_flight_snapshot(InFlightSnapshot { + compaction: Some(protocol::InFlightCompaction::from_running(&lifecycle).unwrap()), + ..InFlightSnapshot::default() + }); + + let mut unrelated = lifecycle; + unrelated.compaction_id = "another-compaction".into(); + unrelated.revision = 2; + unrelated.state = protocol::CompactionLifecycleState::Failed; + unrelated.error = Some("must not replace".into()); + app.handle_worker_event(Event::CompactFailed { + lifecycle: unrelated, + }); + + assert_eq!(compact_block_count(&app), 1); + assert!(matches!( + app.blocks.as_slice(), + [Block::Compact(CompactEvent::Streaming { .. })] + )); + } + #[test] fn shutdown_marks_live_compact_incomplete() { let mut app = App::new("test".into()); diff --git a/crates/worker/src/in_flight.rs b/crates/worker/src/in_flight.rs index 39e1d844..539bead1 100644 --- a/crates/worker/src/in_flight.rs +++ b/crates/worker/src/in_flight.rs @@ -23,6 +23,7 @@ pub(crate) struct InFlightInner { next_block_id: u64, blocks: Vec, commands: Vec, + compaction: Option, } #[derive(Debug, Clone)] @@ -53,6 +54,7 @@ impl InFlightEvents { next_block_id: 1, blocks: Vec::new(), commands: Vec::new(), + compaction: None, })), working_event_tx, } @@ -219,6 +221,35 @@ impl InFlightEvents { self.lock().commands = commands; } + /// Publish current compaction progress into reconnect snapshots. + /// Terminal lifecycle events clear the in-flight value; their durable + /// session record remains the historical authority. + pub(crate) fn update_compaction(&self, lifecycle: &protocol::CompactionLifecycle) { + let mut inner = self.lock(); + match lifecycle.state { + protocol::CompactionLifecycleState::Running => match &inner.compaction { + None => inner.compaction = protocol::InFlightCompaction::from_running(lifecycle), + Some(current) + if current.compaction_id == lifecycle.compaction_id + && lifecycle.revision > current.revision => + { + inner.compaction = protocol::InFlightCompaction::from_running(lifecycle); + } + Some(_) => {} + }, + protocol::CompactionLifecycleState::Done + | protocol::CompactionLifecycleState::Failed + | protocol::CompactionLifecycleState::Interrupted => { + if inner.compaction.as_ref().is_some_and(|current| { + current.compaction_id == lifecycle.compaction_id + && lifecycle.revision > current.revision + }) { + inner.compaction = None; + } + } + } + } + pub(crate) fn clear(&self) { let mut inner = self.lock(); inner.clear(); @@ -378,6 +409,7 @@ impl InFlightInner { .filter_map(TrackedBlock::to_snapshot_block) .collect(), commands: self.commands.clone(), + compaction: self.compaction.clone(), } } @@ -740,6 +772,82 @@ mod tests { assert!(snapshot_from_guard(&guard).commands.is_empty()); } + #[test] + fn compaction_progress_is_snapshot_only_while_running() { + let (working_event_tx, _) = broadcast::channel(16); + let in_flight = InFlightEvents::new(working_event_tx); + let running = protocol::CompactionLifecycle { + schema_version: 3, + compaction_id: "compact-1".into(), + revision: 1, + internal_worker: None, + state: protocol::CompactionLifecycleState::Running, + started_at_ms: 100, + ended_at_ms: None, + summary: None, + error: None, + new_segment_id: None, + }; + + in_flight.update_compaction(&running); + let guard = in_flight.snapshot_guard(); + assert_eq!( + snapshot_from_guard(&guard) + .compaction + .as_ref() + .map(|item| item.compaction_id.as_str()), + Some("compact-1") + ); + assert!(!snapshot_from_guard(&guard).is_empty()); + drop(guard); + + let mut stale = running.clone(); + stale.revision = 0; + in_flight.update_compaction(&stale); + let mut other = running.clone(); + other.compaction_id = "compact-2".into(); + other.revision = 2; + in_flight.update_compaction(&other); + let guard = in_flight.snapshot_guard(); + assert_eq!( + snapshot_from_guard(&guard) + .compaction + .as_ref() + .map(|item| (item.compaction_id.as_str(), item.revision)), + Some(("compact-1", 1)) + ); + drop(guard); + + in_flight.clear(); + let guard = in_flight.snapshot_guard(); + assert_eq!( + snapshot_from_guard(&guard) + .compaction + .as_ref() + .map(|item| item.compaction_id.as_str()), + Some("compact-1") + ); + drop(guard); + + let mut stale_done = running.clone(); + stale_done.state = protocol::CompactionLifecycleState::Done; + in_flight.update_compaction(&stale_done); + let guard = in_flight.snapshot_guard(); + assert!(snapshot_from_guard(&guard).compaction.is_some()); + drop(guard); + + let mut done = running; + done.revision = 2; + done.state = protocol::CompactionLifecycleState::Done; + done.ended_at_ms = Some(200); + done.summary = Some("private summary".into()); + done.new_segment_id = Some("staged-segment".into()); + in_flight.update_compaction(&done); + + let guard = in_flight.snapshot_guard(); + assert!(snapshot_from_guard(&guard).compaction.is_none()); + } + #[test] fn clear_discards_uncommitted_blocks_without_protocol_event() { let (working_event_tx, _) = broadcast::channel(16); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 38414970..97147ac8 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -4567,6 +4567,9 @@ impl Worker { lifecycle: CompactionLifecycle, ) -> Result<(), WorkerError> { self.persist_compaction_lifecycle(&lifecycle)?; + if let Some(in_flight) = &self.in_flight { + in_flight.update_compaction(&lifecycle); + } self.send_event(Event::CompactStart { lifecycle }); Ok(()) } @@ -4576,6 +4579,9 @@ impl Worker { lifecycle: CompactionLifecycle, ) -> Result<(), WorkerError> { self.persist_compaction_lifecycle(&lifecycle)?; + if let Some(in_flight) = &self.in_flight { + in_flight.update_compaction(&lifecycle); + } self.send_event(Event::CompactFailed { lifecycle }); Ok(()) } @@ -5057,6 +5063,9 @@ impl Worker { match outcome { Ok((new_segment_id, _summary)) => { debug_assert_eq!(lifecycle.state, CompactionLifecycleState::Done); + if let Some(in_flight) = &self.in_flight { + in_flight.update_compaction(&lifecycle); + } self.send_event(Event::CompactDone { lifecycle: lifecycle.clone(), }); @@ -5661,54 +5670,66 @@ impl Worker { }); self.store .create_segment(old_loc.session_id, new_segment_id, &initial_entries)?; + + // The name-keyed Worker metadata pointer is the durable commit point for + // Segment activation. Everything above is staging: a failure leaves the + // current Segment and live session untouched, while the unreachable + // replacement may be collected later. The metadata store publishes the + // complete record with an atomic replace, so restore observes either the + // old Segment or this fully-written replacement, never a partial switch. + self.write_worker_metadata_active(SegmentLocation { + session_id: old_loc.session_id, + segment_id: new_segment_id, + })?; + + // All live mutations after the durable commit are infallible and happen + // before the replacement SegmentStart is broadcast. This keeps the + // append destination, Session projection, Engine cache identity, and + // reconnect snapshot on one side of the same activation boundary. self.segment_state.set_location(SegmentLocation { session_id: old_loc.session_id, segment_id: new_segment_id, }); self.segment_state .set_entries_written(initial_entries.len()); - // Broadcast the complete compacted prefix. Runtime-owned extensions - // must remain visible and restorable with the replacement segment. - self.sink - .reset_with_initial_entries(initial_entries.clone()); - // Keep workers.json pointing at the live segment_id. Without this - // a concurrent `restore_from_manifest(new_segment_id)` would - // see no live writer and grab the session this Worker just moved - // into, causing two writers to race on the same jsonl. Skipped - // when no allocation is installed (e.g. compact under - // `Worker::new` in tests). - if self.scope_allocation.is_some() { - worker_allocation::update_segment(&self.manifest.worker.name, new_segment_id)?; - } - self.write_worker_metadata_active(SegmentLocation { - session_id: old_loc.session_id, - segment_id: new_segment_id, - })?; - // Align user_segments with the post-compaction history. Items - // before `retain_from` (now folded into the summary) lose their - // segments; only the user_messages surviving in retained_items - // keep them. They are always the trailing K entries of - // `self.user_segments` because submissions are appended in order. self.user_segments = retained_user_segments; - self.session.replace_history(compacted_history_entries); - // Compaction-introduced system messages are part of the new - // SegmentStart's history (broadcast above) — clients derive - // their blocks from `SegmentStart.history`. No per-item - // broadcast is required. - let _ = &compact_introduced_system_messages; let worker = self.engine.as_mut().unwrap(); - // Anchor the prompt cache at the summary item so that Anthropic - // can place a durable `cache_control` breakpoint there — our - // compact layout guarantees history[0] is the summary. worker.set_cache_anchor(Some(0)); - // Re-key the OpenAI Responses prompt cache namespace to the new - // segment_id so post-compact turns use the rewritten session namespace. worker.set_cache_key(Some(new_segment_id.to_string())); self.usage_history .lock() .expect("usage_history poisoned") .clear(); + + // workers.json is live writer-bookkeeping derived from the durable + // metadata pointer above. A stale value still retains the worker-name + // lease and cannot authorize another writer; after a process exit it is + // reclaimed through the normal stale-allocation path. Do not report a + // committed compaction as failed solely because this derived projection + // could not be refreshed. + if self.scope_allocation.is_some() + && let Err(error) = + worker_allocation::update_segment(&self.manifest.worker.name, new_segment_id) + { + warn!( + worker = %self.manifest.worker.name, + segment_id = %new_segment_id, + error = %error, + "compaction committed but live writer allocation projection could not be refreshed" + ); + } + + // Broadcast only after every live authority points at the committed + // replacement. Runtime-owned extensions stay visible and restorable. + self.sink + .reset_with_initial_entries(initial_entries.clone()); + + // Compaction-introduced system messages are part of the new + // SegmentStart's history (broadcast above) — clients derive + // their blocks from `SegmentStart.history`. No per-item + // broadcast is required. + let _ = &compact_introduced_system_messages; Ok((new_segment_id, summary_text)) } diff --git a/crates/worker/tests/compact_events_test.rs b/crates/worker/tests/compact_events_test.rs index 1521b909..719adafd 100644 --- a/crates/worker/tests/compact_events_test.rs +++ b/crates/worker/tests/compact_events_test.rs @@ -8,7 +8,7 @@ use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use agen::Engine; use agen::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent}; @@ -17,14 +17,103 @@ use agen::llm_client::{ClientError, LlmClient, Request}; use async_trait::async_trait; use futures::Stream; use protocol::{Event, Method, RunResult}; -use session_store::{CombinedStore, FsWorkerStore, WorkerMetadataStore}; -use session_store::{FsStore, LogEntry, Store}; +use session_store::{ + CombinedStore, FsStore, FsWorkerStore, LogEntry, Store, WorkerMetadata, WorkerMetadataStore, + WorkerStoreError, +}; use tokio::sync::broadcast; use worker::{Worker, WorkerController}; type TestStore = CombinedStore; +#[derive(Clone)] +struct FaultingWorkerMetadataStore { + inner: FsWorkerStore, + fail_next_update: Arc, +} + +impl FaultingWorkerMetadataStore { + fn new(root: impl Into) -> Self { + Self { + inner: FsWorkerStore::new(root).unwrap(), + fail_next_update: Arc::new(AtomicBool::new(false)), + } + } + + fn arm_update_failure(&self) { + self.fail_next_update.store(true, Ordering::SeqCst); + } +} + +impl WorkerMetadataStore for FaultingWorkerMetadataStore { + fn write(&self, metadata: &WorkerMetadata) -> Result<(), WorkerStoreError> { + let old_segment_id = self + .inner + .read_by_name(&metadata.worker_name)? + .and_then(|current| current.active) + .and_then(|active| active.segment_id); + let new_segment_id = metadata + .active + .as_ref() + .and_then(|active| active.segment_id); + if old_segment_id != new_segment_id && self.fail_next_update.swap(false, Ordering::SeqCst) { + return Err(WorkerStoreError::Io(std::io::Error::other( + "injected active Segment commit failure", + ))); + } + self.inner.write(metadata) + } + + fn read_by_name(&self, worker_name: &str) -> Result, WorkerStoreError> { + self.inner.read_by_name(worker_name) + } + + fn update_by_name( + &self, + worker_name: &str, + mutate: F, + ) -> Result + where + F: FnOnce(&mut WorkerMetadata), + { + let mut metadata = self + .inner + .read_by_name(worker_name)? + .unwrap_or_else(|| WorkerMetadata::new(worker_name, None)); + let old_segment_id = metadata + .active + .as_ref() + .and_then(|active| active.segment_id); + mutate(&mut metadata); + let new_segment_id = metadata + .active + .as_ref() + .and_then(|active| active.segment_id); + if old_segment_id != new_segment_id && self.fail_next_update.swap(false, Ordering::SeqCst) { + return Err(WorkerStoreError::Io(std::io::Error::other( + "injected active Segment commit failure", + ))); + } + self.inner.write(&metadata)?; + Ok(metadata) + } + + fn list_names(&self) -> Result, WorkerStoreError> { + self.inner.list_names() + } + + fn root_dir(&self) -> Option { + self.inner.root_dir() + } + + fn delete_by_name(&self, worker_name: &str) -> Result<(), WorkerStoreError> { + self.inner.delete_by_name(worker_name) + } +} + +type FaultingTestStore = CombinedStore; + fn annotated(item: Item) -> session_store::LoggedHistoryEntry { session_store::LoggedHistoryEntry { item: session_store::LoggedItem::from(item), @@ -229,6 +318,41 @@ async fn make_worker(client: MockClient) -> Worker { make_worker_with_manifest(POST_RUN_MANIFEST_TOML, client).await } +async fn make_faulting_worker( + client: MockClient, +) -> ( + Worker, + FaultingWorkerMetadataStore, + FsStore, +) { + let manifest = worker::WorkerManifest::from_toml(MID_TURN_MANIFEST_TOML).unwrap(); + let store_tmp = tempfile::tempdir().unwrap(); + let segment_store = FsStore::new(store_tmp.path()).unwrap(); + let metadata_store = FaultingWorkerMetadataStore::new(store_tmp.path().join("pods")); + let store = CombinedStore::new(segment_store.clone(), metadata_store.clone()); + std::mem::forget(store_tmp); + + let pwd_tmp = tempfile::tempdir().unwrap(); + let pwd = pwd_tmp.path().to_path_buf(); + let scope = worker::Scope::writable(&pwd).unwrap(); + std::mem::forget(pwd_tmp); + + let engine = + Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client); + let mut worker = Worker::new( + manifest, + engine, + store, + worker::WorkerWorkspaceContext::local_filesystem(None), + worker::WorkerFilesystemAuthority::local(pwd.clone(), pwd.clone()), + scope, + ) + .await + .unwrap(); + worker.enable_worker_metadata_write_through().unwrap(); + (worker, metadata_store, segment_store) +} + /// Drain whatever events are already queued on `rx`. Non-blocking. fn drain(rx: &mut broadcast::Receiver) -> Vec { let mut out = Vec::new(); @@ -282,6 +406,45 @@ fn system_texts_in_sink_session_start( Vec::new() } +#[tokio::test] +async fn failed_active_segment_commit_keeps_live_and_durable_history_on_old_segment() { + let client = MockClient::new(vec![ + single_text_events("seed response"), + write_summary_tool_use_events("summary-1", "replacement summary"), + single_text_events("continued on old segment"), + ]); + let (mut worker, metadata_store, segment_store) = make_faulting_worker(client).await; + worker.run_text("seed input").await.unwrap(); + let old_segment_id = worker.segment_id(); + + metadata_store.arm_update_failure(); + let error = worker.compact(0).await.unwrap_err(); + assert!( + error + .to_string() + .contains("injected active Segment commit failure") + ); + + assert_eq!(worker.segment_id(), old_segment_id); + let metadata = metadata_store + .read_by_name("test-worker") + .unwrap() + .expect("active Worker metadata should remain present"); + assert_eq!( + metadata.active.and_then(|active| active.segment_id), + Some(old_segment_id) + ); + + worker.run_text("continue input").await.unwrap(); + let active_records = segment_store + .read_all(worker.session_id(), old_segment_id) + .unwrap(); + assert!( + format!("{active_records:?}").contains("continue input"), + "the live Worker must continue appending to the old active Segment" + ); +} + /// Worker metadata starts with a reserved Session and no Segment, then becomes /// active once the first SegmentStart is materialized by `run`. #[tokio::test] diff --git a/web/workspace/src/lib/generated/protocol.ts b/web/workspace/src/lib/generated/protocol.ts index 10b573b0..2b1ba468 100644 --- a/web/workspace/src/lib/generated/protocol.ts +++ b/web/workspace/src/lib/generated/protocol.ts @@ -113,7 +113,16 @@ export type RewindSummary = { truncated_to_entries: number, discarded_entries: n export type InFlightBlock = { "kind": "text", text: string, finished?: boolean, } | { "kind": "thinking", text: string, finished?: boolean, } | { "kind": "tool_call", id: string, name: string, args: string, state?: InFlightToolCallState, }; -export type InFlightSnapshot = { blocks?: Array, commands?: Array, }; +export type InFlightCompaction = { schema_version: number, compaction_id: string, revision: number, internal_worker: InternalWorkerRef | null, started_at_ms: number, }; + +export type InFlightSnapshot = { blocks?: Array, commands?: Array, +/** + * The currently running compaction, if any. + * + * This is lifecycle progress only. Candidate history and the staged + * Segment remain private until the Segment is activated atomically. + */ +compaction?: InFlightCompaction | null, }; export type SessionEntryProvenance = "human_input" | "worker_input" | "flow_instruction" | "backend_instruction" | "model_output" | "tool_output" | "derived_summary" | "legacy_unknown"; diff --git a/web/workspace/src/lib/workspace/console/model.test.ts b/web/workspace/src/lib/workspace/console/model.test.ts index 4303bec2..e2c29985 100644 --- a/web/workspace/src/lib/workspace/console/model.test.ts +++ b/web/workspace/src/lib/workspace/console/model.test.ts @@ -1079,6 +1079,27 @@ Deno.test("projectConsole upserts compaction lifecycle by stable id", () => { assertEquals(completed.lines[0].streaming, false); }); +Deno.test("snapshot restores running compaction without staged content", () => { + const snapshot = snapshotEvent("/repo") as Extract; + snapshot.data.in_flight = { + blocks: [], + compaction: { + schema_version: 3, + compaction_id: "compaction-snapshot", + revision: 1, + internal_worker: null, + started_at_ms: 1_000, + }, + }; + + const projection = projectConsole([{ eventId: "snapshot", event: snapshot }]); + + assertEquals(projection.lines.length, 1); + assertEquals(projection.lines[0].id, "compaction-compaction-snapshot"); + assertEquals(projection.lines[0].compaction?.state, "running"); + assertEquals(projection.lines[0].body.includes("staged"), false); +}); + Deno.test("compaction service activity stays nested in one lifecycle item", () => { const worker = { session_id: "compactor-session", diff --git a/web/workspace/src/lib/workspace/console/model.ts b/web/workspace/src/lib/workspace/console/model.ts index 890db3ed..e44dcc0a 100644 --- a/web/workspace/src/lib/workspace/console/model.ts +++ b/web/workspace/src/lib/workspace/console/model.ts @@ -6,6 +6,7 @@ import type { CompactionLifecycle, Event as ProtocolEvent, InFlightBlock, + InFlightCompaction, InFlightToolCallState, InternalWorkerRef, InternalWorkerSnapshot, @@ -638,7 +639,7 @@ function projectInternalWorkerSnapshot( eventId: string, cwd: string | null, ): InternalWorkerProjection { - const console = snapshotProjectionFromSession( + let console = snapshotProjectionFromSession( `${eventId}:internal:${snapshot.worker.session_id}:snapshot`, snapshot.session, cwd, @@ -669,6 +670,9 @@ function projectInternalWorkerSnapshot( console.internalWorkers = (snapshot.internal_workers ?? []).map((child) => projectInternalWorkerSnapshot(child, eventId, cwd) ); + if (snapshot.in_flight?.compaction) { + console = applyInFlightCompaction(console, snapshot.in_flight.compaction); + } return { worker: snapshot.worker, revision: snapshot.revision, console }; } @@ -714,6 +718,20 @@ function compactionActivity( .filter((value, index, values) => value.length > 0 && values.indexOf(value) === index); } +function applyInFlightCompaction( + projection: ConsoleProjection, + progress: InFlightCompaction, +): ConsoleProjection { + return applyCompactionLifecycle(projection, { + ...progress, + state: "running", + ended_at_ms: null, + summary: null, + error: null, + new_segment_id: null, + }); +} + function applyCompactionLifecycle( projection: ConsoleProjection, lifecycle: CompactionLifecycle @@ -970,6 +988,13 @@ export function applyProtocolEvent( projectInternalWorkerSnapshot(worker, envelope.eventId, next.cwd) ); next.removedInternalWorkers = {}; + if (event.data.in_flight?.compaction) { + const withCompaction = applyInFlightCompaction( + next, + event.data.in_flight.compaction, + ); + next.lines = withCompaction.lines; + } for (const line of next.lines) { const compaction = line.compaction; if (!compaction) continue; @@ -981,19 +1006,6 @@ export function applyProtocolEvent( activity: compactionActivity(next, sessionId), }; } - if ( - compaction.state === "running" && - (!sessionId || !next.internalWorkers.some((worker) => - worker.worker.session_id === sessionId - )) - ) { - line.streaming = false; - line.compaction = { - ...line.compaction!, - state: "interrupted", - endedAtMs: envelope.observedAtMs ?? Date.now(), - }; - } } applyWorkerStateSnapshot(next, event.data.state, envelope.eventId); break;