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

fix: make compaction activation atomic

Browse Source
This commit is contained in:
Hare
2026-09-16 03:44:48 +09:00
parent 016dbd7cb1
commit 4e7a314a00
9 changed files with 523 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/protocol/src/lib.rs
+54 -2
View File
@@ -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<InternalWorkerRef>,
pub started_at_ms: u64,
}
impl InFlightCompaction {
pub fn from_running(lifecycle: &CompactionLifecycle) -> Option<Self> {
(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<InFlightBlock>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub commands: Vec<CommandSnapshot>,
/// 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<InFlightCompaction>,
}
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::<Event>(&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:?}"),
}
}
crates/protocol/src/typescript.rs
+10 -9
View File
@@ -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::<RewindTarget>(&cfg, &mut output);
push_decl::<RewindSummary>(&cfg, &mut output);
push_decl::<InFlightBlock>(&cfg, &mut output);
push_decl::<InFlightCompaction>(&cfg, &mut output);
push_decl::<InFlightSnapshot>(&cfg, &mut output);
push_decl::<SessionEntryProvenance>(&cfg, &mut output);
push_decl::<SessionMessageRole>(&cfg, &mut output);
crates/tui/src/app.rs
+70 -1
View File
@@ -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<String>,
/// 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<InternalWorkerView>,
@@ -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 { .. } => {
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());
crates/worker/src/in_flight.rs
+108
View File
@@ -23,6 +23,7 @@ pub(crate) struct InFlightInner {
next_block_id: u64,
blocks: Vec<TrackedBlock>,
commands: Vec<CommandSnapshot>,
compaction: Option<protocol::InFlightCompaction>,
}
#[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);
crates/worker/src/worker.rs
+54 -33
View File
@@ -4567,6 +4567,9 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
});
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))
}
crates/worker/tests/compact_events_test.rs
+166 -3
View File
@@ -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<FsStore, FsWorkerStore>;
#[derive(Clone)]
struct FaultingWorkerMetadataStore {
inner: FsWorkerStore,
fail_next_update: Arc<AtomicBool>,
}
impl FaultingWorkerMetadataStore {
fn new(root: impl Into<std::path::PathBuf>) -> 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<Option<WorkerMetadata>, WorkerStoreError> {
self.inner.read_by_name(worker_name)
}
fn update_by_name<F>(
&self,
worker_name: &str,
mutate: F,
) -> Result<WorkerMetadata, WorkerStoreError>
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<Vec<String>, WorkerStoreError> {
self.inner.list_names()
}
fn root_dir(&self) -> Option<std::path::PathBuf> {
self.inner.root_dir()
}
fn delete_by_name(&self, worker_name: &str) -> Result<(), WorkerStoreError> {
self.inner.delete_by_name(worker_name)
}
}
type FaultingTestStore = CombinedStore<FsStore, FaultingWorkerMetadataStore>;
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<MockClient, TestStore> {
make_worker_with_manifest(POST_RUN_MANIFEST_TOML, client).await
}
async fn make_faulting_worker(
client: MockClient,
) -> (
Worker<MockClient, FaultingTestStore>,
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<Event>) -> Vec<Event> {
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]
web/workspace/src/lib/generated/protocol.ts
+10 -1
View File
@@ -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<InFlightBlock>, commands?: Array<CommandSnapshot>, };
export type InFlightCompaction = { schema_version: number, compaction_id: string, revision: number, internal_worker: InternalWorkerRef | null, started_at_ms: number, };
export type InFlightSnapshot = { blocks?: Array<InFlightBlock>, commands?: Array<CommandSnapshot>,
/**
* 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";
web/workspace/src/lib/workspace/console/model.test.ts
+21
View File
@@ -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<Event, { event: "snapshot" }>;
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",
web/workspace/src/lib/workspace/console/model.ts
+26 -14
View File
@@ -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;