Make compaction transitions recoverable

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