fix: make compaction activation atomic

crates/worker/src/in_flight.rs

@@ -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

@@ -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

@@ -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]