From 2a14b93335cac9841bdc8cfee6f143b25947db2c Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Sun, 26 Jul 2026 10:56:19 +0200 Subject: [PATCH] Make compaction transitions recoverable --- src/agent.rs | 29 +++++++++++++++ src/app.rs | 55 ++++++++++++++++++++++++++++ src/app/generation.rs | 56 +++++++++++++++++++++++++--- src/app/projects.rs | 11 +++++- src/app/view.rs | 4 ++ src/database.rs | 19 +++++++++- src/engine.rs | 67 ++++++++++++++++++++++++++++++---- src/engine/gguf.rs | 27 ++++++++++++++ src/engine/metal.rs | 5 ++- src/engine/metal/checkpoint.rs | 8 ++++ src/engine/tokenizer.rs | 42 ++++++++++++++++----- src/engine/validation.rs | 25 +++++++++++++ src/runtime.rs | 7 +++- 13 files changed, 329 insertions(+), 26 deletions(-) diff --git a/src/agent.rs b/src/agent.rs index 89cba3a..c64766e 100644 --- a/src/agent.rs +++ b/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 { + 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::>(); + 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 { 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!( diff --git a/src/app.rs b/src/app.rs index d100468..ba2b6f8 100644 --- a/src/app.rs +++ b/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 { diff --git a/src/app/generation.rs b/src/app/generation.rs index 7c4245f..28da7f5 100644 --- a/src/app/generation.rs +++ b/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::>(); + 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(()) } diff --git a/src/app/projects.rs b/src/app/projects.rs index 46da271..5b55872 100644 --- a/src/app/projects.rs +++ b/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 diff --git a/src/app/view.rs b/src/app/view.rs index 7332d5c..d107cd0 100644 --- a/src/app/view.rs +++ b/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 { diff --git a/src/database.rs b/src/database.rs index a27dc16..d161404 100644 --- a/src/database.rs +++ b/src/database.rs @@ -407,13 +407,24 @@ impl Database { session_id: i32, summary: &str, tail: &[MessageDraft], + context_used: u32, + context_limit: u32, ) -> Result, 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::), + )) .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()); diff --git a/src/engine.rs b/src/engine.rs index 5df4924..abd16a2 100644 --- a/src/engine.rs +++ b/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 { 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 { + 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, 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), ) -> Result { + 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 - .executor - .model() - .render_conversation(&rebuilt_system, &tail, settings.reasoning_mode) - .len() as u32; + let history_tokens = self.executor.model().render_history( + &rebuilt_system, + &tail, + settings.reasoning_mode, + ); + 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(), }) })(); - self.executor.reset()?; - self.checkpoint = None; + 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 { 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 diff --git a/src/engine/gguf.rs b/src/engine/gguf.rs index 68341dc..e39cc19 100644 --- a/src/engine/gguf.rs +++ b/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::>(); + 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() } diff --git a/src/engine/metal.rs b/src/engine/metal.rs index 0808805..2ac3d6e 100644 --- a/src/engine/metal.rs +++ b/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, }) diff --git a/src/engine/metal/checkpoint.rs b/src/engine/metal/checkpoint.rs index 945d42e..7d04a3d 100644 --- a/src/engine/metal/checkpoint.rs +++ b/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 diff --git a/src/engine/tokenizer.rs b/src/engine/tokenizer.rs index 9964381..9fe0489 100644 --- a/src/engine/tokenizer.rs +++ b/src/engine/tokenizer.rs @@ -210,6 +210,25 @@ impl Tokenizer { system_prompt: &str, messages: &[ChatTurn], reasoning: ReasoningMode, + ) -> Vec { + self.encode_messages(system_prompt, messages, reasoning, true) + } + + pub(super) fn encode_history( + &self, + system_prompt: &str, + messages: &[ChatTurn], + reasoning: ReasoningMode, + ) -> Vec { + self.encode_messages(system_prompt, messages, reasoning, false) + } + + fn encode_messages( + &self, + system_prompt: &str, + messages: &[ChatTurn], + reasoning: ReasoningMode, + continue_assistant: bool, ) -> Vec { 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,17 +300,22 @@ 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); } } - output.push(self.assistant); - if reasoning != ReasoningMode::Direct { - output.push(self.think_start); - } else if self.family == ModelFamily::Glm { - output.extend([self.think_start, self.think_end]); - } else { - output.push(self.think_end); + if continue_assistant { + output.push(self.assistant); + if reasoning != ReasoningMode::Direct { + output.push(self.think_start); + } else if self.family == ModelFamily::Glm { + output.extend([self.think_start, self.think_end]); + } else { + output.push(self.think_end); + } } output } diff --git a/src/engine/validation.rs b/src/engine/validation.rs index 7b4daaa..5dd8d33 100644 --- a/src/engine/validation.rs +++ b/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, diff --git a/src/runtime.rs b/src/runtime.rs index 202255c..74f9eb7 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -122,6 +122,7 @@ impl GenerationService { turn: TurnSettings, messages: Vec, reason: &str, + checkpoint: PathBuf, idle_timeout: Duration, ) -> Result { 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, );