hugo/DS4Server
1
0
Fork 0
Code Issues 8 Pull Requests Actions Packages Projects Releases Wiki Activity

Make compaction transitions recoverable

Browse Source
This commit is contained in:
Georg Bauer
2026-07-26 10:56:19 +02:00
parent 6aa45b2cf0
commit 2a14b93335
13 changed files with 329 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/agent.rs
View File

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

55
src/app.rs
View File

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

56
src/app/generation.rs
View File

@@ -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(())
}

11
src/app/projects.rs
View File

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

4
src/app/view.rs
View File

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

19
src/database.rs
View File

@@ -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());

67
src/engine.rs
View File

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

27
src/engine/gguf.rs
View File

@@ -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()
}

5
src/engine/metal.rs
View File

@@ -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,
})

8
src/engine/metal/checkpoint.rs
View File

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

42
src/engine/tokenizer.rs
View File

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

25
src/engine/validation.rs
View File

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

7
src/runtime.rs
View File

@@ -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,
);