Reuse canonical chat context

src/engine.rs

@@ -560,6 +560,7 @@ impl Generator {
     pub(crate) fn generate(
         &mut self,
         checkpoint: &Path,
+        bootstrap: Option<&Path>,
         messages: &[ChatTurn],
         settings: &TurnSettings,
         cancelled: &AtomicBool,
@@ -578,6 +579,11 @@ impl Generator {
         if checkpoint_present && !selected.found {
             phase(checkpoint_rebuild_activity(selected.incompatible));
         }
+        if !selected.found
+            && let Some(directory) = bootstrap
+        {
+            self.prepare_bootstrap(directory, settings, cancelled, &mut progress, &mut phase)?;
+        }
         let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
         self.publish_execution_stats();
         let (mut output, prompt_complete) = result?;
@@ -596,6 +602,8 @@ impl Generator {
         output.checkpoint_bytes = std::fs::metadata(checkpoint)
             .map(|item| item.len())
             .unwrap_or(0);
+        self.checkpoint = Some(checkpoint.to_owned());
+        self.resident_active = Some(checkpoint.to_owned());
         Ok(output)
     }
 
@@ -648,32 +656,8 @@ impl Generator {
         let history_key =
             conversation_key(&settings.system_prompt, settings.reasoning_mode, history);
         let history_tag: [u8; 32] = Sha256::digest(&history_key).into();
-        let mut previous_checkpoint = self.checkpoint.clone();
-        if self.executor.checkpoint_tag() != history_tag {
-            if let Some(entry) = store.find(&history_key, self.executor.context()) {
-                if self.select_checkpoint(&entry.checkpoint, entry.tag)?.found {
-                    store.touch(&entry)?;
-                    self.last_store_tokens = entry.tokens;
-                    previous_checkpoint = Some(entry.checkpoint);
-                } else {
-                    store.discard(&entry);
-                    self.last_store_tokens = 0;
-                    previous_checkpoint = None;
-                }
-            } else {
-                let key = resident_key(directory, history_tag);
-                let restored = self.activate_resident(key)?;
-                if !restored || self.executor.checkpoint_tag() != history_tag {
-                    self.executor.reset()?;
-                }
-                self.checkpoint = None;
-                self.last_store_tokens = 0;
-                previous_checkpoint = None;
-                self.metrics.kv_lookup(KvLookup::Miss);
-            }
-        } else {
-            self.metrics.kv_lookup(KvLookup::MemoryHit);
-        }
+        let previous_checkpoint =
+            self.restore_cached_prefix(directory, &store, &history_key, history_tag)?;
 
         let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
         self.publish_execution_stats();
@@ -731,6 +715,98 @@ impl Generator {
         Ok(output)
     }
 
+    fn prepare_bootstrap(
+        &mut self,
+        directory: &Path,
+        settings: &TurnSettings,
+        cancelled: &AtomicBool,
+        progress: &mut impl FnMut(u32, u32, Option<f32>),
+        phase: &mut impl FnMut(&'static str),
+    ) -> Result<(), String> {
+        let key = conversation_key(&settings.system_prompt, settings.reasoning_mode, &[]);
+        let tag: [u8; 32] = Sha256::digest(&key).into();
+        let store = KvStore::open(directory, settings.kv_cache.budget_bytes)?;
+        self.restore_cached_prefix(directory, &store, &key, tag)?;
+
+        let tokens = self.executor.model().render_history(
+            &settings.system_prompt,
+            &[],
+            settings.reasoning_mode,
+        );
+        if tokens.len() >= self.executor.context() as usize {
+            return Err(format!(
+                "System prompt has {} tokens, but the configured context size is {} tokens",
+                tokens.len(),
+                self.executor.context()
+            ));
+        }
+        let reused = self.executor.align_prompt(&tokens)?;
+        if reused == tokens.len() {
+            return Ok(());
+        }
+
+        phase("Updating system prompt cache…");
+        let completed = self.prefill_suffix(&tokens, reused, cancelled, progress)?;
+        if completed != tokens.len() - reused {
+            return Err("generation cancelled while updating the system prompt cache".into());
+        }
+        self.executor.note_checkpoint_tag(tag);
+        if !settings.kv_cache.stores(self.executor.position(), true, 0) {
+            self.checkpoint = None;
+            self.resident_active = Some(resident_key(directory, tag));
+            return Ok(());
+        }
+
+        let checkpoint = store.checkpoint_path(&key);
+        self.save_checkpoint(&checkpoint, tag)?;
+        let retained = store.record(
+            &checkpoint,
+            &key,
+            tag,
+            self.executor.position(),
+            self.executor.context(),
+            StoreReason::Cold,
+        )?;
+        self.checkpoint = retained.then_some(checkpoint.clone());
+        self.resident_active = Some(if retained {
+            checkpoint
+        } else {
+            resident_key(directory, tag)
+        });
+        Ok(())
+    }
+
+    fn restore_cached_prefix(
+        &mut self,
+        directory: &Path,
+        store: &KvStore,
+        key: &[u8],
+        tag: [u8; 32],
+    ) -> Result<Option<PathBuf>, String> {
+        if self.executor.checkpoint_tag() == tag {
+            self.metrics.kv_lookup(KvLookup::MemoryHit);
+            return Ok(self.checkpoint.clone());
+        }
+        if let Some(entry) = store.find(key, self.executor.context()) {
+            if self.select_checkpoint(&entry.checkpoint, entry.tag)?.found {
+                store.touch(&entry)?;
+                self.last_store_tokens = entry.tokens;
+                return Ok(Some(entry.checkpoint));
+            }
+            store.discard(&entry);
+        } else {
+            let key = resident_key(directory, tag);
+            let restored = self.activate_resident(key)?;
+            if !restored || self.executor.checkpoint_tag() != tag {
+                self.executor.reset()?;
+            }
+            self.metrics.kv_lookup(KvLookup::Miss);
+        }
+        self.checkpoint = None;
+        self.last_store_tokens = 0;
+        Ok(None)
+    }
+
     #[allow(clippy::too_many_arguments)]
     pub(crate) fn compact(
         &mut self,
@@ -986,6 +1062,34 @@ impl Generator {
         result
     }
 
+    fn prefill_suffix(
+        &mut self,
+        tokens: &[i32],
+        reused: usize,
+        cancelled: &AtomicBool,
+        progress: &mut impl FnMut(u32, u32, Option<f32>),
+    ) -> Result<usize, String> {
+        let suffix = &tokens[reused..];
+        if (reused == 0 && tokens.len() > 1) || suffix.len() >= 4 {
+            let context = self.executor.context();
+            self.executor.prefill(suffix, |used| {
+                progress(used, context, None);
+                !cancelled.load(Ordering::Relaxed)
+            })
+        } else {
+            let mut completed = 0;
+            for &token in suffix {
+                if cancelled.load(Ordering::Relaxed) {
+                    break;
+                }
+                self.executor.eval(token)?;
+                completed += 1;
+                progress(self.executor.position(), self.executor.context(), None);
+            }
+            Ok(completed)
+        }
+    }
+
     fn generate_inner(
         &mut self,
         messages: &[ChatTurn],
@@ -1047,24 +1151,7 @@ impl Generator {
         let mut pending_utf8 = Vec::new();
         let prompt_tokens = tokens.len();
         let suffix = &tokens[reused..];
-        let completed = if (reused == 0 && tokens.len() > 1) || suffix.len() >= 4 {
-            let context = self.executor.context();
-            self.executor.prefill(suffix, |used| {
-                progress(used, context, None);
-                !cancelled.load(Ordering::Relaxed)
-            })?
-        } else {
-            let mut completed = 0;
-            for &token in suffix {
-                if cancelled.load(Ordering::Relaxed) {
-                    break;
-                }
-                self.executor.eval(token)?;
-                completed += 1;
-                progress(self.executor.position(), self.executor.context(), None);
-            }
-            completed
-        };
+        let completed = self.prefill_suffix(&tokens, reused, cancelled, progress)?;
         self.publish_execution_stats();
         if completed != suffix.len() {
             return Ok((
@@ -1648,6 +1735,34 @@ mod sampling_tests {
         assert!(conversation_key("System", ReasoningMode::High, &messages).starts_with(&prefix));
     }
 
+    #[test]
+    fn bootstrap_key_is_the_prefix_before_dynamic_session_context() {
+        let system = "System\n\nProject instructions from AGENTS.md:\n\nkeep this";
+        let bootstrap = conversation_key(system, ReasoningMode::High, &[]);
+        let messages = vec![
+            ChatTurn {
+                user: false,
+                tool: false,
+                system: true,
+                skip_previous_eos: false,
+                reasoning: None,
+                reasoning_complete: true,
+                content: "current date and time".into(),
+            },
+            ChatTurn {
+                user: true,
+                tool: false,
+                system: false,
+                skip_previous_eos: false,
+                reasoning: None,
+                reasoning_complete: true,
+                content: "hello".into(),
+            },
+        ];
+
+        assert!(conversation_key(system, ReasoningMode::High, &messages).starts_with(&bootstrap));
+    }
+
     #[test]
     fn checkpoint_rebuilds_explain_compatibility_and_history_misses() {
         assert!(checkpoint_rebuild_activity(true).contains("different model"));