Make compaction transitions recoverable

src/agent.rs

@@ -276,6 +276,32 @@ impl Tools {
         (self.context_tokens.max(4096) as usize * 2).min(512 * 1024)
     }
 
+    pub(crate) fn compaction_observation(&mut self) -> Option<String> {
+        let mut running = self
+            .jobs
+            .values_mut()
+            .filter_map(|job| {
+                job.child.try_wait().ok().flatten().is_none().then(|| {
+                    format!(
+                        "bash job={} pid={} status=running command={}\noutput_path={}\n",
+                        job.id,
+                        job.child.id(),
+                        job.command,
+                        job.output.display()
+                    )
+                })
+            })
+            .collect::<Vec<_>>();
+        if running.is_empty() {
+            return None;
+        }
+        running.sort();
+        Some(format!(
+            "Bash job update after context compaction. Running jobs still need explicit bash_status or bash_stop if relevant.\n{}",
+            running.concat()
+        ))
+    }
+
     fn existing_path(&self, value: &str) -> Result<PathBuf, String> {
         let path = if Path::new(value).is_absolute() {
             PathBuf::from(value)
@@ -1288,6 +1314,9 @@ mod tests {
             ],
         );
         assert!(tools.execute(&running, &cancel).contains("status=running"));
+        let observation = tools.compaction_observation().unwrap();
+        assert!(observation.contains("bash job=2"));
+        assert!(observation.contains("status=running"));
         let stopped = tools.execute(&call("bash_stop", [("job", "2")]), &cancel);
         assert!(!stopped.contains("status=running"));
         assert_eq!(

src/app.rs

@@ -212,6 +212,8 @@ pub(crate) enum Message {
     SessionTitleChanged(String),
     ConfirmRenameSession,
     RetitleSession(i32),
+    CompactSession(i32),
+    RebuildSessionContext(i32),
     SetSessionState(i32, SessionState),
     ToggleArchivedSessions(i32),
     ShowChat,
@@ -773,6 +775,10 @@ impl App {
                 if let Some(active) = &self.active_tools {
                     active.cancel.store(true, Ordering::Relaxed);
                 }
+                #[cfg(target_os = "macos")]
+                if let Some(compaction) = &self.active_compaction {
+                    compaction.active.cancel.store(true, Ordering::Relaxed);
+                }
             }
             Message::GenerationTick => {
                 #[cfg(target_os = "macos")]
@@ -924,6 +930,44 @@ impl App {
                     self.error = Some("Local Metal generation requires macOS.".into());
                 }
             }
+            Message::CompactSession(session_id) => {
+                self.session_menu = None;
+                if self.generating || self.selected_session != Some(session_id) {
+                    self.error = Some("Open an idle session before compacting it.".into());
+                } else {
+                    #[cfg(target_os = "macos")]
+                    if let Err(error) = self.start_compaction(
+                        generation::PendingContinuation::None,
+                        "manual compact action",
+                    ) {
+                        self.error = Some(error);
+                    }
+                    #[cfg(not(target_os = "macos"))]
+                    {
+                        self.error = Some("Local Metal generation requires macOS.".into());
+                    }
+                }
+            }
+            Message::RebuildSessionContext(session_id) => {
+                self.session_menu = None;
+                if self.generating && self.selected_session == Some(session_id) {
+                    self.error =
+                        Some("Stop the active generation before rebuilding context.".into());
+                } else {
+                    match fs::remove_file(session_checkpoint_path(session_id)) {
+                        Ok(()) => {
+                            self.error = None;
+                            self.finish_cache_change();
+                        }
+                        Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
+                            self.error = None;
+                        }
+                        Err(error) => {
+                            self.error = Some(format!("Could not discard the checkpoint: {error}"));
+                        }
+                    }
+                }
+            }
             Message::SetSessionState(session_id, state) => {
                 self.session_menu = None;
                 if let Some(database) = &mut self.database {
@@ -1402,6 +1446,13 @@ fn sweep_orphan_checkpoints(directory: &Path, projects: &[ProjectWithSessions])
     let mut removed = false;
     for file in files.flatten() {
         let path = file.path();
+        if path
+            .extension()
+            .is_some_and(|value| value == "compacting" || value == "tmp")
+        {
+            removed |= fs::remove_file(&path).is_ok();
+            continue;
+        }
         if path.extension().is_none_or(|value| value != "bin") {
             continue;
         }
@@ -1427,6 +1478,10 @@ fn session_checkpoint_path(session_id: i32) -> PathBuf {
     kv_cache_path().join(format!("{session_id}.bin"))
 }
 
+fn session_compaction_checkpoint_path(session_id: i32) -> PathBuf {
+    kv_cache_path().join(format!("{session_id}.compacting"))
+}
+
 /// Content-addressed KV cache for turns that belong to no stored session: the
 /// HTTP endpoint and the app's own one-shot requests share it.
 fn transient_cache_path() -> PathBuf {

src/app/generation.rs

@@ -20,13 +20,14 @@ pub(super) struct TitleRequest {
 
 #[cfg(target_os = "macos")]
 pub(super) enum PendingContinuation {
+    None,
     User(String),
     Tool(String),
 }
 
 #[cfg(target_os = "macos")]
 pub(super) struct CompactionRequest {
-    active: ActiveGeneration,
+    pub(super) active: ActiveGeneration,
     pending: PendingContinuation,
 }
 
@@ -92,6 +93,16 @@ impl App {
             return;
         }
         #[cfg(target_os = "macos")]
+        if prompt == "/compact" {
+            self.composer.clear();
+            if let Err(error) =
+                self.start_compaction(PendingContinuation::None, "manual /compact request")
+            {
+                self.error = Some(error);
+            }
+            return;
+        }
+        #[cfg(target_os = "macos")]
         if !std::mem::take(&mut self.skip_compaction_once)
             && crate::compaction::should_compact(self.context_used, self.context_limit)
         {
@@ -466,7 +477,7 @@ impl App {
     }
 
     #[cfg(target_os = "macos")]
-    fn start_compaction(
+    pub(super) fn start_compaction(
         &mut self,
         pending: PendingContinuation,
         reason: &str,
@@ -506,6 +517,10 @@ impl App {
                 effective.turn,
                 messages,
                 reason,
+                session_compaction_checkpoint_path(
+                    self.selected_session
+                        .ok_or_else(|| "The active session is unavailable.".to_owned())?,
+                ),
                 idle_timeout,
             )?;
         self.active_compaction = Some(CompactionRequest { active, pending });
@@ -537,6 +552,7 @@ impl App {
                     match result {
                         Ok(compacted) => {
                             if let Err(error) = self.apply_compaction(&compacted) {
+                                let _ = fs::remove_file(&compacted.checkpoint);
                                 self.generating = false;
                                 self.activity = None;
                                 self.error = Some(error);
@@ -546,6 +562,7 @@ impl App {
                             self.generating = false;
                             self.activity = None;
                             match request.pending {
+                                PendingContinuation::None => {}
                                 PendingContinuation::User(prompt) => {
                                     self.composer = prompt;
                                     self.skip_compaction_once = true;
@@ -595,7 +612,7 @@ impl App {
         let session_id = self
             .selected_session
             .ok_or_else(|| "The active session is unavailable.".to_owned())?;
-        let tail = compacted
+        let mut tail = compacted
             .tail
             .iter()
             .map(|message| crate::database::MessageDraft {
@@ -606,11 +623,32 @@ impl App {
                 content: message.content.clone(),
             })
             .collect::<Vec<_>>();
+        if let Some((tools_session, tools)) = &self.agent_tools
+            && *tools_session == session_id
+            && let Some(observation) = tools
+                .lock()
+                .unwrap_or_else(|poisoned| poisoned.into_inner())
+                .compaction_observation()
+        {
+            tail.push(crate::database::MessageDraft {
+                user: false,
+                tool: true,
+                reasoning: None,
+                reasoning_complete: true,
+                content: observation,
+            });
+        }
         let messages = self
             .database
             .as_mut()
             .ok_or_else(|| "The project database is unavailable.".to_owned())?
-            .replace_with_compacted_transcript(session_id, &compacted.summary, &tail)
+            .replace_with_compacted_transcript(
+                session_id,
+                &compacted.summary,
+                &tail,
+                compacted.context_tokens,
+                self.context_limit,
+            )
             .map_err(|error| format!("Could not save compacted conversation: {error}"))?;
         self.conversation = messages.into_iter().map(ChatMessage::from).collect();
         if let Some(session) = self
@@ -620,8 +658,16 @@ impl App {
             .find(|session| session.id == session_id)
         {
             session.compacted_summary = Some(compacted.summary.clone());
+            session.context_used = compacted.context_tokens as i32;
+            session.context_limit = self.context_limit as i32;
+            session.last_tokens_per_second = None;
         }
-        let _ = fs::remove_file(session_checkpoint_path(session_id));
+        let final_checkpoint = session_checkpoint_path(session_id);
+        fs::rename(&compacted.checkpoint, &final_checkpoint).map_err(|error| {
+            format!(
+                "The compacted conversation was saved, but its checkpoint could not be promoted: {error}. It will rebuild on next use."
+            )
+        })?;
         Ok(())
     }
 

src/app/projects.rs

@@ -196,7 +196,14 @@ mod tests {
         let directory =
             std::env::temp_dir().join(format!("ds4-server-sweep-{}", std::process::id()));
         std::fs::create_dir_all(&directory).unwrap();
-        for name in ["1.bin", "2.bin", "notes.bin", "7.bin"] {
+        for name in [
+            "1.bin",
+            "2.bin",
+            "notes.bin",
+            "7.bin",
+            "1.compacting",
+            "2.tmp",
+        ] {
             std::fs::write(directory.join(name), b"payload").unwrap();
         }
         let projects = vec![ProjectWithSessions {
@@ -213,6 +220,8 @@ mod tests {
         // Not a session checkpoint, so not ours to delete.
         assert!(directory.join("notes.bin").exists());
         assert!(!directory.join("7.bin").exists());
+        assert!(!directory.join("1.compacting").exists());
+        assert!(!directory.join("2.tmp").exists());
         // Nothing left to sweep on the next pass.
         assert!(!super::super::sweep_orphan_checkpoints(
             &directory, &projects

src/app/view.rs

@@ -446,6 +446,10 @@ impl App {
                 .on_press(Message::StartRenameSession(session.id)),
             menu_action(ICON_SPARK, "Retitle with AI")
                 .on_press(Message::RetitleSession(session.id)),
+            menu_action(ICON_SPARK, "Compact context")
+                .on_press(Message::CompactSession(session.id)),
+            menu_action(ICON_ARCHIVE, "Rebuild context on next use")
+                .on_press(Message::RebuildSessionContext(session.id)),
         ]
         .spacing(4);
         actions = match state {

src/database.rs

@@ -407,13 +407,24 @@ impl Database {
         session_id: i32,
         summary: &str,
         tail: &[MessageDraft],
+        context_used: u32,
+        context_limit: u32,
     ) -> Result<Vec<StoredMessage>, String> {
+        let context_used = i32::try_from(context_used)
+            .map_err(|_| "Compacted context is too large to save".to_owned())?;
+        let context_limit = i32::try_from(context_limit)
+            .map_err(|_| "Context limit is too large to save".to_owned())?;
         self.connection
             .transaction(|connection| {
                 diesel::delete(messages::table.filter(messages::session_id.eq(session_id)))
                     .execute(connection)?;
                 diesel::update(sessions::table.find(session_id))
-                    .set(sessions::compacted_summary.eq(Some(summary)))
+                    .set((
+                        sessions::compacted_summary.eq(Some(summary)),
+                        sessions::context_used.eq(context_used),
+                        sessions::context_limit.eq(context_limit),
+                        sessions::last_tokens_per_second.eq(None::<f32>),
+                    ))
                     .execute(connection)?;
                 let mut stored = Vec::with_capacity(tail.len());
                 for message in tail {
@@ -561,6 +572,8 @@ mod tests {
                     reasoning_complete: true,
                     content: "Recent question".into(),
                 }],
+                321,
+                32_768,
             )
             .unwrap();
         assert_eq!(compacted.len(), 1);
@@ -572,6 +585,10 @@ mod tests {
                 .as_deref(),
             Some("Keep the active task.")
         );
+        assert_eq!(
+            reopened.load_projects().unwrap()[0].sessions[0].context_used,
+            321
+        );
         assert_eq!(reopened.load_messages(session.id).unwrap().len(), 1);
         reopened.delete_session(session.id).unwrap();
         assert!(reopened.load_messages(session.id).unwrap().is_empty());

src/engine.rs

@@ -247,6 +247,16 @@ impl Model {
         }
     }
 
+    fn checkpoint_identity(&self) -> [u8; 32] {
+        let mut hash = Sha256::new();
+        hash.update(b"DS4Server model checkpoint identity v1");
+        hash.update(self.main.checkpoint_identity());
+        if let Some(support) = &self.support {
+            hash.update(support.checkpoint_identity());
+        }
+        hash.finalize().into()
+    }
+
     pub(crate) fn tokenize(&self, text: &str) -> Vec<i32> {
         self.tokenizer.tokenize(text)
     }
@@ -270,6 +280,15 @@ impl Model {
             .encode_conversation(system, messages, reasoning)
     }
 
+    fn render_history(
+        &self,
+        system: &str,
+        messages: &[ChatTurn],
+        reasoning: ReasoningMode,
+    ) -> Vec<i32> {
+        self.tokenizer.encode_history(system, messages, reasoning)
+    }
+
     fn render_continuation(
         &self,
         prompt: &str,
@@ -346,6 +365,7 @@ pub(crate) struct CompactionOutput {
     pub(crate) summary: String,
     pub(crate) tail: Vec<ChatTurn>,
     pub(crate) context_tokens: u32,
+    pub(crate) checkpoint: PathBuf,
 }
 
 #[cfg(target_os = "macos")]
@@ -508,9 +528,11 @@ impl Generator {
         messages: &[ChatTurn],
         settings: &TurnSettings,
         reason: &str,
+        checkpoint: &Path,
         cancelled: &AtomicBool,
         mut progress: impl FnMut(u32, u32, Option<f32>),
     ) -> Result<CompactionOutput, String> {
+        let _ = std::fs::remove_file(checkpoint);
         self.executor.reset()?;
         self.checkpoint = None;
         let result = (|| {
@@ -549,6 +571,10 @@ impl Generator {
                 return Err("context compaction produced an empty summary".into());
             }
 
+            // The private request must never become the rebuilt session prefix.
+            self.executor.reset()?;
+            self.checkpoint = None;
+
             let full = self.executor.model().render_conversation(
                 &settings.system_prompt,
                 messages,
@@ -576,19 +602,38 @@ impl Generator {
             let tail = messages[start..].to_vec();
             let rebuilt_system =
                 crate::compaction::summary_system_prompt(&settings.system_prompt, Some(&summary));
-            let context_tokens = self
+            let history_tokens = self.executor.model().render_history(
-                .executor
+                &rebuilt_system,
-                .model()
+                &tail,
-                .render_conversation(&rebuilt_system, &tail, settings.reasoning_mode)
+                settings.reasoning_mode,
-                .len() as u32;
+            );
+            if history_tokens.len() >= self.executor.context() as usize {
+                return Err("compacted context does not fit the configured context".into());
+            }
+            let context = self.executor.context();
+            let completed = self.executor.prefill(&history_tokens, |used| {
+                progress(used, context, None);
+                !cancelled.load(Ordering::Relaxed)
+            })?;
+            if completed != history_tokens.len() || cancelled.load(Ordering::Relaxed) {
+                return Err("context compaction interrupted during rebuild".into());
+            }
+            let tag = conversation_tag(&rebuilt_system, settings.reasoning_mode, &tail);
+            self.save_checkpoint(checkpoint, tag)?;
             Ok(CompactionOutput {
                 summary,
                 tail,
-                context_tokens,
+                context_tokens: history_tokens.len() as u32,
+                checkpoint: checkpoint.to_owned(),
             })
         })();
+        if result.is_err() {
+            let _ = std::fs::remove_file(checkpoint);
             self.executor.reset()?;
             self.checkpoint = None;
+        } else {
+            self.checkpoint = Some(checkpoint.to_owned());
+        }
         result
     }
 
@@ -598,6 +643,12 @@ impl Generator {
         expected_tag: [u8; 32],
     ) -> Result<bool, String> {
         if self.checkpoint.as_deref() == Some(checkpoint) {
+            if !checkpoint.is_file() {
+                self.executor.reset()?;
+                self.checkpoint = None;
+                self.metrics.kv_lookup(KvLookup::Miss);
+                return Ok(false);
+            }
             let found = self.executor.checkpoint_tag() == expected_tag;
             self.metrics.kv_lookup(if found {
                 KvLookup::MemoryHit

src/engine/gguf.rs

@@ -1,4 +1,5 @@
 use memmap2::{Mmap, MmapOptions};
+use sha2::{Digest, Sha256};
 use std::collections::HashMap;
 use std::fs::File;
 use std::path::{Path, PathBuf};
@@ -176,6 +177,32 @@ impl Gguf {
         self.map.len() as u64
     }
 
+    pub(super) fn checkpoint_identity(&self) -> [u8; 32] {
+        let mut hash = Sha256::new();
+        hash.update(b"DS4Server GGUF checkpoint identity v1");
+        hash.update(
+            self.path
+                .canonicalize()
+                .unwrap_or_else(|_| self.path.clone())
+                .to_string_lossy()
+                .as_bytes(),
+        );
+        hash.update(self.len().to_le_bytes());
+        hash.update(self.data_offset.to_le_bytes());
+        let mut tensors = self.tensors.iter().collect::<Vec<_>>();
+        tensors.sort_by_key(|(name, _)| *name);
+        for (name, tensor) in tensors {
+            hash.update(name.as_bytes());
+            hash.update(tensor.kind.to_le_bytes());
+            hash.update(tensor.offset.to_le_bytes());
+            hash.update(tensor.bytes.to_le_bytes());
+            for dimension in &tensor.dims {
+                hash.update(dimension.to_le_bytes());
+            }
+        }
+        hash.finalize().into()
+    }
+
     pub(super) fn map_ptr(&self) -> *const u8 {
         self.map.as_ptr()
     }

src/engine/metal.rs

@@ -14,7 +14,7 @@ use std::ptr::NonNull;
 use std::time::UNIX_EPOCH;
 
 const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4RKV01";
-const CHECKPOINT_VERSION: u32 = 1;
+const CHECKPOINT_VERSION: u32 = 2;
 const CHECKPOINT_IO_CHUNK: usize = 8 * 1024 * 1024;
 const DEFAULT_PREFILL_CHUNK: u32 = 4096;
 
@@ -571,6 +571,7 @@ pub(super) struct Executor {
     quality: bool,
     checkpoint_tag: [u8; 32],
     model_modified: (u64, u32),
+    model_identity: [u8; 32],
     _context: Context,
     model: Model,
 }
@@ -599,6 +600,7 @@ impl Executor {
             .and_then(|modified| modified.duration_since(UNIX_EPOCH).ok())
             .map(|duration| (duration.as_secs(), duration.subsec_nanos()))
             .unwrap_or_default();
+        let model_identity = model.checkpoint_identity();
         Ok(Self {
             weights,
             session,
@@ -607,6 +609,7 @@ impl Executor {
             quality,
             checkpoint_tag: [0; 32],
             model_modified,
+            model_identity,
             _context: context_handle,
             model,
         })

src/engine/metal/checkpoint.rs

@@ -81,6 +81,8 @@ impl Executor {
         write_u64(file, self.model.main.len())?;
         write_u64(file, self.model_modified.0)?;
         write_u32(file, self.model_modified.1)?;
+        file.write_all(&self.model_identity)
+            .map_err(|error| error.to_string())?;
         for weight in [self.weights.token_embedding, self.weights.output] {
             write_u64(file, weight.offset)?;
             write_u64(file, weight.bytes)?;
@@ -185,6 +187,12 @@ impl Executor {
         if read_u64(file)? != self.model_modified.0 || read_u32(file)? != self.model_modified.1 {
             return Err("KV checkpoint model file has changed".into());
         }
+        let mut model_identity = [0; 32];
+        file.read_exact(&mut model_identity)
+            .map_err(|error| error.to_string())?;
+        if model_identity != self.model_identity {
+            return Err("KV checkpoint model identity or quantization changed".into());
+        }
         for weight in [self.weights.token_embedding, self.weights.output] {
             if read_u64(file)? != weight.offset
                 || read_u64(file)? != weight.bytes

src/engine/tokenizer.rs

@@ -210,6 +210,25 @@ impl Tokenizer {
         system_prompt: &str,
         messages: &[ChatTurn],
         reasoning: ReasoningMode,
+    ) -> Vec<i32> {
+        self.encode_messages(system_prompt, messages, reasoning, true)
+    }
+
+    pub(super) fn encode_history(
+        &self,
+        system_prompt: &str,
+        messages: &[ChatTurn],
+        reasoning: ReasoningMode,
+    ) -> Vec<i32> {
+        self.encode_messages(system_prompt, messages, reasoning, false)
+    }
+
+    fn encode_messages(
+        &self,
+        system_prompt: &str,
+        messages: &[ChatTurn],
+        reasoning: ReasoningMode,
+        continue_assistant: bool,
     ) -> Vec<i32> {
         let mut output = vec![self.bos];
         if self.family == ModelFamily::Glm && self.sop >= 0 {
@@ -235,7 +254,7 @@ impl Tokenizer {
             }
             output.extend(self.tokenize_rendered(system_prompt));
         }
-        for message in messages {
+        for (index, message) in messages.iter().enumerate() {
             if message.tool {
                 if self.family == ModelFamily::Glm {
                     output.push(self.observation);
@@ -281,10 +300,14 @@ impl Tokenizer {
                 }
             }
             output.extend(self.tokenize_rendered(&message.content));
-            if !message.user && self.family == ModelFamily::DeepSeek {
+            if !message.user
+                && self.family == ModelFamily::DeepSeek
+                && (continue_assistant || index + 1 < messages.len())
+            {
                 output.push(self.eos);
             }
         }
+        if continue_assistant {
             output.push(self.assistant);
             if reasoning != ReasoningMode::Direct {
                 output.push(self.think_start);
@@ -293,6 +316,7 @@ impl Tokenizer {
             } else {
                 output.push(self.think_end);
             }
+        }
         output
     }
 

src/engine/validation.rs

@@ -941,6 +941,31 @@ mod tests {
             model.render_continuation("Hello", ReasoningMode::Direct, true),
             [128_803, 19_923, 128_804, 128_822]
         );
+        assert_eq!(
+            model.render_history(
+                "",
+                &[
+                    ChatTurn {
+                        user: true,
+                        tool: false,
+                        skip_previous_eos: false,
+                        reasoning: None,
+                        reasoning_complete: true,
+                        content: "Hello".into(),
+                    },
+                    ChatTurn {
+                        user: false,
+                        tool: false,
+                        skip_previous_eos: false,
+                        reasoning: None,
+                        reasoning_complete: true,
+                        content: "Hello".into(),
+                    },
+                ],
+                ReasoningMode::Direct,
+            ),
+            [0, 128_803, 19_923, 128_804, 128_822, 19_923]
+        );
         let tool_turn = ChatTurn {
             user: false,
             tool: true,

src/runtime.rs

@@ -122,6 +122,7 @@ impl GenerationService {
         turn: TurnSettings,
         messages: Vec<ChatTurn>,
         reason: &str,
+        checkpoint: PathBuf,
         idle_timeout: Duration,
     ) -> Result<ActiveGeneration, String> {
         let cancel = Arc::new(AtomicBool::new(false));
@@ -132,7 +133,7 @@ impl GenerationService {
                 engine,
                 turn,
                 messages,
-                checkpoint: CheckpointTarget::OneShot(PathBuf::new()),
+                checkpoint: CheckpointTarget::Local(checkpoint),
                 compact_reason: Some(reason.to_owned()),
                 idle_timeout,
                 cancel: Arc::clone(&cancel),
@@ -261,10 +262,14 @@ fn run_command(
             });
         };
         if let Some(reason) = &command.compact_reason {
+            let CheckpointTarget::Local(checkpoint) = &command.checkpoint else {
+                unreachable!("compaction checkpoints are local")
+            };
             let result = generator.compact(
                 &command.messages,
                 &command.turn,
                 reason,
+                checkpoint,
                 &command.cancel,
                 &mut progress,
             );