Make compaction transitions recoverable

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

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,