Add prefix-aware KV session cache

src/engine.rs

@@ -1,5 +1,7 @@
 mod gguf;
 #[cfg(target_os = "macos")]
+mod kvstore;
+#[cfg(target_os = "macos")]
 mod metal;
 mod tokenizer;
 mod validation;
@@ -10,6 +12,8 @@ use crate::model::ModelChoice;
 use crate::settings::TurnSettings;
 use crate::settings::{EngineSettings, ReasoningMode};
 use gguf::{F16, F32, Gguf, I32, IQ2_XXS, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Tensor, Value};
+#[cfg(target_os = "macos")]
+use kvstore::{KvStore, StoreReason};
 use sha2::{Digest, Sha256};
 use std::path::{Path, PathBuf};
 #[cfg(target_os = "macos")]
@@ -409,15 +413,25 @@ impl Generator {
         let history = messages
             .split_last()
             .map_or(messages, |(_, history)| history);
-        let history_tag =
-            conversation_tag(&settings.system_prompt, settings.reasoning_mode, history);
-        let checkpoint = directory.join(format!("{}.bin", hex_tag(history_tag)));
+        let store = KvStore::open(directory)?;
+        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 {
-            self.select_checkpoint(&checkpoint, history_tag)?;
-            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)? {
+                    store.touch(&entry)?;
+                    previous_checkpoint = Some(entry.checkpoint);
+                } else {
+                    store.discard(&entry);
+                    previous_checkpoint = None;
+                }
+            } else {
                 self.executor.reset()?;
                 self.checkpoint = None;
-                let _ = std::fs::remove_file(&checkpoint);
+                previous_checkpoint = None;
+                self.metrics.kv_lookup(KvLookup::Miss);
             }
         } else {
             self.metrics.kv_lookup(KvLookup::MemoryHit);
@@ -427,30 +441,34 @@ impl Generator {
             self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?;
         let mut completed = messages.to_vec();
         completed.push(output.message.clone());
-        let completed_tag =
-            conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed);
-        std::fs::create_dir_all(directory).map_err(|error| {
-            format!(
-                "Could not create transient KV cache directory {}: {error}",
-                directory.display()
-            )
-        })?;
-        let completed_checkpoint = directory.join(format!("{}.bin", hex_tag(completed_tag)));
-        output.previous_checkpoint_bytes = std::fs::metadata(&completed_checkpoint)
-            .ok()
+        output.previous_checkpoint_bytes = previous_checkpoint
+            .as_deref()
+            .and_then(|path| std::fs::metadata(path).ok())
             .map(|item| item.len());
-        self.save_checkpoint(
-            &completed_checkpoint,
-            if prompt_complete {
-                completed_tag
-            } else {
-                [0; 32]
-            },
-        )?;
-        output.checkpoint_bytes = std::fs::metadata(&completed_checkpoint)
-            .map(|item| item.len())
-            .unwrap_or(0);
-        self.checkpoint = Some(completed_checkpoint);
+        if prompt_complete {
+            let completed_key =
+                conversation_key(&settings.system_prompt, settings.reasoning_mode, &completed);
+            let completed_tag: [u8; 32] = Sha256::digest(&completed_key).into();
+            let completed_checkpoint = store.checkpoint_path(&completed_key);
+            self.save_checkpoint(&completed_checkpoint, completed_tag)?;
+            let retained = store.record(
+                &completed_checkpoint,
+                &completed_key,
+                completed_tag,
+                self.executor.position(),
+                self.executor.context(),
+                if history.is_empty() {
+                    StoreReason::Cold
+                } else {
+                    StoreReason::Continued
+                },
+            )?;
+            output.checkpoint_bytes = retained
+                .then(|| std::fs::metadata(&completed_checkpoint).ok())
+                .flatten()
+                .map_or(0, |item| item.len());
+            self.checkpoint = retained.then_some(completed_checkpoint);
+        }
         Ok(output)
     }
 
@@ -458,15 +476,15 @@ impl Generator {
         &mut self,
         checkpoint: &Path,
         expected_tag: [u8; 32],
-    ) -> Result<(), String> {
+    ) -> Result<bool, String> {
         if self.checkpoint.as_deref() == Some(checkpoint) {
-            self.metrics
-                .kv_lookup(if self.executor.checkpoint_tag() == expected_tag {
-                    KvLookup::MemoryHit
-                } else {
-                    KvLookup::Miss
-                });
-            return Ok(());
+            let found = self.executor.checkpoint_tag() == expected_tag;
+            self.metrics.kv_lookup(if found {
+                KvLookup::MemoryHit
+            } else {
+                KvLookup::Miss
+            });
+            return Ok(found);
         }
         self.executor.reset()?;
         self.metrics.kv_read_started();
@@ -476,8 +494,9 @@ impl Generator {
             .load_checkpoint(checkpoint, &mut |bytes| self.metrics.kv_read_bytes(bytes));
         self.metrics
             .kv_read_finished(started.elapsed(), loaded.is_err());
+        let found = matches!(loaded, Ok(true)) && self.executor.checkpoint_tag() == expected_tag;
         let lookup = match loaded {
-            Ok(true) if self.executor.checkpoint_tag() == expected_tag => KvLookup::DiskHit,
+            Ok(true) if found => KvLookup::DiskHit,
             Ok(true) | Ok(false) => KvLookup::Miss,
             Err(_) => {
                 self.executor.reset()?;
@@ -487,7 +506,7 @@ impl Generator {
         };
         self.metrics.kv_lookup(lookup);
         self.checkpoint = Some(checkpoint.to_owned());
-        Ok(())
+        Ok(found)
     }
 
     fn save_checkpoint(&mut self, checkpoint: &Path, tag: [u8; 32]) -> Result<(), String> {
@@ -871,44 +890,38 @@ fn emit_safe_text(
 }
 
 #[cfg(target_os = "macos")]
-fn hex_tag(tag: [u8; 32]) -> String {
-    const HEX: &[u8; 16] = b"0123456789abcdef";
-    let mut output = String::with_capacity(64);
-    for byte in tag {
-        output.push(HEX[(byte >> 4) as usize] as char);
-        output.push(HEX[(byte & 0x0f) as usize] as char);
+fn conversation_key(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn]) -> Vec<u8> {
+    fn text(output: &mut Vec<u8>, value: &str) {
+        output.extend_from_slice(&(value.len() as u64).to_le_bytes());
+        output.extend_from_slice(value.as_bytes());
+    }
+
+    let mut output = b"DS4Server chat checkpoint v2".to_vec();
+    text(&mut output, system);
+    output.push(match reasoning {
+        ReasoningMode::Direct => 0,
+        ReasoningMode::High => 1,
+        ReasoningMode::Max => 2,
+    });
+    for message in messages {
+        output.push(u8::from(message.user));
+        output.push(u8::from(message.skip_previous_eos));
+        match &message.reasoning {
+            Some(reasoning) => {
+                output.push(1);
+                text(&mut output, reasoning);
+            }
+            None => output.push(0),
+        }
+        output.push(u8::from(message.reasoning_complete));
+        text(&mut output, &message.content);
     }
     output
 }
 
 #[cfg(target_os = "macos")]
 fn conversation_tag(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn]) -> [u8; 32] {
-    fn text(hasher: &mut Sha256, value: &str) {
-        hasher.update((value.len() as u64).to_le_bytes());
-        hasher.update(value.as_bytes());
-    }
-
-    let mut hasher = Sha256::new();
-    hasher.update(b"DS4Server chat checkpoint v1");
-    text(&mut hasher, system);
-    hasher.update([match reasoning {
-        ReasoningMode::Direct => 0,
-        ReasoningMode::High => 1,
-        ReasoningMode::Max => 2,
-    }]);
-    for message in messages {
-        hasher.update([u8::from(message.user)]);
-        match &message.reasoning {
-            Some(reasoning) => {
-                hasher.update([1]);
-                text(&mut hasher, reasoning);
-            }
-            None => hasher.update([0]),
-        }
-        hasher.update([u8::from(message.reasoning_complete)]);
-        text(&mut hasher, &message.content);
-    }
-    hasher.finalize().into()
+    Sha256::digest(conversation_key(system, reasoning, messages)).into()
 }
 
 #[cfg(target_os = "macos")]
@@ -1061,7 +1074,7 @@ mod sampling_tests {
 
     #[test]
     fn checkpoint_tag_covers_the_canonical_chat_state() {
-        let messages = [ChatTurn {
+        let mut messages = vec![ChatTurn {
             user: true,
             skip_previous_eos: false,
             reasoning: None,
@@ -1081,6 +1094,20 @@ mod sampling_tests {
             tag,
             conversation_tag("System", ReasoningMode::Direct, &messages)
         );
+        messages[0].skip_previous_eos = true;
+        assert_ne!(
+            tag,
+            conversation_tag("System", ReasoningMode::High, &messages)
+        );
+        let prefix = conversation_key("System", ReasoningMode::High, &messages);
+        messages.push(ChatTurn {
+            user: false,
+            skip_previous_eos: false,
+            reasoning: Some("because".into()),
+            reasoning_complete: true,
+            content: "Hi".into(),
+        });
+        assert!(conversation_key("System", ReasoningMode::High, &messages).starts_with(&prefix));
     }
 
     #[test]

src/engine/kvstore.rs

@@ -0,0 +1,385 @@
+use sha2::{Digest, Sha256};
+use std::cmp::Ordering;
+use std::fs::{self, File};
+use std::io::{Read, Write};
+use std::path::{Path, PathBuf};
+use std::time::{SystemTime, UNIX_EPOCH};
+
+const INDEX_MAGIC: &[u8; 8] = b"DS4KVIDX";
+const INDEX_VERSION: u32 = 1;
+const DEFAULT_BUDGET_BYTES: u64 = 4 * 1024 * 1024 * 1024;
+const MAX_KEY_BYTES: u64 = 256 * 1024 * 1024;
+const HIT_HALF_LIFE_SECONDS: f64 = 6.0 * 60.0 * 60.0;
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub(super) enum StoreReason {
+    Cold,
+    Continued,
+}
+
+impl StoreReason {
+    fn code(self) -> u8 {
+        match self {
+            Self::Cold => 1,
+            Self::Continued => 2,
+        }
+    }
+
+    fn from_code(value: u8) -> Option<Self> {
+        Some(match value {
+            1 => Self::Cold,
+            2 => Self::Continued,
+            _ => return None,
+        })
+    }
+
+    fn is_anchor(self) -> bool {
+        self == Self::Cold
+    }
+}
+
+#[derive(Clone, Debug)]
+pub(super) struct Entry {
+    pub(super) checkpoint: PathBuf,
+    pub(super) tag: [u8; 32],
+    pub(super) tokens: u32,
+    context: u32,
+    hits: u32,
+    created_at: u64,
+    last_used: u64,
+    reason: StoreReason,
+    key: Vec<u8>,
+    bytes: u64,
+}
+
+pub(super) struct KvStore {
+    directory: PathBuf,
+    budget_bytes: u64,
+}
+
+impl KvStore {
+    pub(super) fn open(directory: &Path) -> Result<Self, String> {
+        fs::create_dir_all(directory).map_err(|error| {
+            format!(
+                "Could not create transient KV cache directory {}: {error}",
+                directory.display()
+            )
+        })?;
+        Ok(Self {
+            directory: directory.to_owned(),
+            budget_bytes: DEFAULT_BUDGET_BYTES,
+        })
+    }
+
+    #[cfg(test)]
+    fn with_budget(directory: &Path, budget_bytes: u64) -> Result<Self, String> {
+        let mut store = Self::open(directory)?;
+        store.budget_bytes = budget_bytes;
+        Ok(store)
+    }
+
+    pub(super) fn find(&self, prompt_key: &[u8], context: u32) -> Option<Entry> {
+        self.entries()
+            .into_iter()
+            .filter(|entry| entry.context == context && prompt_key.starts_with(&entry.key))
+            .max_by(|left, right| {
+                left.key
+                    .len()
+                    .cmp(&right.key.len())
+                    .then_with(|| left.tokens.cmp(&right.tokens))
+            })
+    }
+
+    pub(super) fn checkpoint_path(&self, key: &[u8]) -> PathBuf {
+        self.directory.join(format!("{}.bin", key_hash(key)))
+    }
+
+    pub(super) fn record(
+        &self,
+        checkpoint: &Path,
+        key: &[u8],
+        tag: [u8; 32],
+        tokens: u32,
+        context: u32,
+        reason: StoreReason,
+    ) -> Result<bool, String> {
+        let checkpoint_bytes = fs::metadata(checkpoint)
+            .map_err(|error| error.to_string())?
+            .len();
+        let index_bytes = index_size(key)?;
+        if checkpoint_bytes.saturating_add(index_bytes) > self.budget_bytes {
+            let _ = fs::remove_file(checkpoint);
+            return Ok(false);
+        }
+        let now = unix_time();
+        let entry = Entry {
+            checkpoint: checkpoint.to_owned(),
+            tag,
+            tokens,
+            context,
+            hits: 0,
+            created_at: now,
+            last_used: now,
+            reason,
+            key: key.to_vec(),
+            bytes: checkpoint_bytes.saturating_add(index_bytes),
+        };
+        write_entry(&entry)?;
+        self.evict(Some(checkpoint));
+        Ok(checkpoint.exists())
+    }
+
+    pub(super) fn touch(&self, entry: &Entry) -> Result<(), String> {
+        let mut touched = entry.clone();
+        touched.hits = touched.hits.saturating_add(1);
+        touched.last_used = unix_time();
+        write_entry(&touched)
+    }
+
+    pub(super) fn discard(&self, entry: &Entry) {
+        let _ = fs::remove_file(&entry.checkpoint);
+        let _ = fs::remove_file(index_path(&entry.checkpoint));
+    }
+
+    fn entries(&self) -> Vec<Entry> {
+        let Ok(files) = fs::read_dir(&self.directory) else {
+            return Vec::new();
+        };
+        files
+            .filter_map(Result::ok)
+            .map(|file| file.path())
+            .filter(|path| path.extension().is_some_and(|value| value == "meta"))
+            .filter_map(|path| read_entry(&path).ok())
+            .filter(|entry| entry.checkpoint.is_file())
+            .collect()
+    }
+
+    fn evict(&self, protected: Option<&Path>) {
+        let mut entries = self.entries();
+        let mut total = entries.iter().map(|entry| entry.bytes).sum::<u64>();
+        let now = unix_time();
+        while total > self.budget_bytes {
+            let Some((index, _)) = entries
+                .iter()
+                .enumerate()
+                .filter(|(_, entry)| protected != Some(entry.checkpoint.as_path()))
+                .min_by(|(_, left), (_, right)| {
+                    eviction_score(left, now)
+                        .partial_cmp(&eviction_score(right, now))
+                        .unwrap_or(Ordering::Equal)
+                        .then_with(|| left.last_used.cmp(&right.last_used))
+                })
+            else {
+                break;
+            };
+            let entry = entries.swap_remove(index);
+            total = total.saturating_sub(entry.bytes);
+            self.discard(&entry);
+        }
+    }
+}
+
+fn eviction_score(entry: &Entry, now: u64) -> f64 {
+    let age = now.saturating_sub(entry.last_used.max(entry.created_at)) as f64;
+    let hits = f64::from(entry.hits) * f64::exp2(-age / HIT_HALF_LIFE_SECONDS);
+    let anchor = if entry.reason.is_anchor() { 2.0 } else { 1.0 };
+    anchor * (hits + 1.0) * f64::from(entry.tokens) / entry.bytes.max(1) as f64
+}
+
+fn write_entry(entry: &Entry) -> Result<(), String> {
+    let path = index_path(&entry.checkpoint);
+    let temporary = path.with_extension("meta.tmp");
+    let mut file = File::create(&temporary).map_err(|error| error.to_string())?;
+    file.write_all(INDEX_MAGIC)
+        .map_err(|error| error.to_string())?;
+    write_u32(&mut file, INDEX_VERSION)?;
+    file.write_all(&entry.tag)
+        .map_err(|error| error.to_string())?;
+    write_u32(&mut file, entry.tokens)?;
+    write_u32(&mut file, entry.context)?;
+    write_u32(&mut file, entry.hits)?;
+    file.write_all(&[entry.reason.code(), 0, 0, 0])
+        .map_err(|error| error.to_string())?;
+    write_u64(&mut file, entry.created_at)?;
+    write_u64(&mut file, entry.last_used)?;
+    write_u64(
+        &mut file,
+        u64::try_from(entry.key.len()).map_err(|_| "KV cache key is too large")?,
+    )?;
+    file.write_all(&entry.key)
+        .map_err(|error| error.to_string())?;
+    file.sync_all().map_err(|error| error.to_string())?;
+    fs::rename(temporary, path).map_err(|error| error.to_string())
+}
+
+fn read_entry(path: &Path) -> Result<Entry, String> {
+    let mut file = File::open(path).map_err(|error| error.to_string())?;
+    let mut magic = [0; 8];
+    file.read_exact(&mut magic)
+        .map_err(|error| error.to_string())?;
+    if &magic != INDEX_MAGIC || read_u32(&mut file)? != INDEX_VERSION {
+        return Err("KV cache index has an invalid signature".into());
+    }
+    let mut tag = [0; 32];
+    file.read_exact(&mut tag)
+        .map_err(|error| error.to_string())?;
+    let tokens = read_u32(&mut file)?;
+    let context = read_u32(&mut file)?;
+    let hits = read_u32(&mut file)?;
+    let mut reason = [0; 4];
+    file.read_exact(&mut reason)
+        .map_err(|error| error.to_string())?;
+    let reason = StoreReason::from_code(reason[0])
+        .ok_or_else(|| "KV cache index has an invalid reason".to_owned())?;
+    let created_at = read_u64(&mut file)?;
+    let last_used = read_u64(&mut file)?;
+    let key_bytes = read_u64(&mut file)?;
+    if tokens == 0 || context == 0 || key_bytes > MAX_KEY_BYTES {
+        return Err("KV cache index has invalid dimensions".into());
+    }
+    let mut key = vec![0; usize::try_from(key_bytes).map_err(|error| error.to_string())?];
+    file.read_exact(&mut key)
+        .map_err(|error| error.to_string())?;
+    let checkpoint = path.with_extension("bin");
+    if checkpoint.file_stem().and_then(|value| value.to_str()) != Some(&key_hash(&key)) {
+        return Err("KV cache index key does not match its filename".into());
+    }
+    let bytes = fs::metadata(&checkpoint)
+        .map_err(|error| error.to_string())?
+        .len()
+        .saturating_add(fs::metadata(path).map_err(|error| error.to_string())?.len());
+    Ok(Entry {
+        checkpoint,
+        tag,
+        tokens,
+        context,
+        hits,
+        created_at,
+        last_used,
+        reason,
+        key,
+        bytes,
+    })
+}
+
+fn index_path(checkpoint: &Path) -> PathBuf {
+    checkpoint.with_extension("meta")
+}
+
+fn index_size(key: &[u8]) -> Result<u64, String> {
+    84_u64
+        .checked_add(u64::try_from(key.len()).map_err(|error| error.to_string())?)
+        .ok_or_else(|| "KV cache index is too large".to_owned())
+}
+
+fn key_hash(key: &[u8]) -> String {
+    let digest: [u8; 32] = Sha256::digest(key).into();
+    const HEX: &[u8; 16] = b"0123456789abcdef";
+    let mut output = String::with_capacity(64);
+    for byte in digest {
+        output.push(HEX[(byte >> 4) as usize] as char);
+        output.push(HEX[(byte & 0x0f) as usize] as char);
+    }
+    output
+}
+
+fn unix_time() -> u64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map_or(0, |duration| duration.as_secs())
+}
+
+fn write_u32(file: &mut File, value: u32) -> Result<(), String> {
+    file.write_all(&value.to_le_bytes())
+        .map_err(|error| error.to_string())
+}
+
+fn write_u64(file: &mut File, value: u64) -> Result<(), String> {
+    file.write_all(&value.to_le_bytes())
+        .map_err(|error| error.to_string())
+}
+
+fn read_u32(file: &mut File) -> Result<u32, String> {
+    let mut bytes = [0; 4];
+    file.read_exact(&mut bytes)
+        .map_err(|error| error.to_string())?;
+    Ok(u32::from_le_bytes(bytes))
+}
+
+fn read_u64(file: &mut File) -> Result<u64, String> {
+    let mut bytes = [0; 8];
+    file.read_exact(&mut bytes)
+        .map_err(|error| error.to_string())?;
+    Ok(u64::from_le_bytes(bytes))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn temporary_directory(name: &str) -> PathBuf {
+        let path = std::env::temp_dir().join(format!(
+            "ds4-server-kvstore-{name}-{}-{}",
+            std::process::id(),
+            unix_time()
+        ));
+        fs::create_dir_all(&path).unwrap();
+        path
+    }
+
+    fn add(store: &KvStore, key: &[u8], tokens: u32, bytes: usize) -> Entry {
+        let path = store.checkpoint_path(key);
+        fs::write(&path, vec![0; bytes]).unwrap();
+        assert!(
+            store
+                .record(
+                    &path,
+                    key,
+                    [tokens as u8; 32],
+                    tokens,
+                    4096,
+                    StoreReason::Continued,
+                )
+                .unwrap()
+        );
+        store.find(key, 4096).unwrap()
+    }
+
+    #[test]
+    fn longest_compatible_prefix_wins_and_hits_survive_reopen() {
+        let directory = temporary_directory("prefix");
+        let store = KvStore::open(&directory).unwrap();
+        add(&store, b"system", 100, 32);
+        let expected = add(&store, b"system-user-one", 200, 32);
+
+        let found = store.find(b"system-user-one-more", 4096).unwrap();
+        assert_eq!(found.checkpoint, expected.checkpoint);
+        store.touch(&found).unwrap();
+        assert_eq!(
+            KvStore::open(&directory)
+                .unwrap()
+                .find(b"system-user-one-more", 4096)
+                .unwrap()
+                .hits,
+            1
+        );
+        assert!(store.find(b"system-user-one-more", 2048).is_none());
+
+        fs::remove_dir_all(directory).unwrap();
+    }
+
+    #[test]
+    fn budget_evicts_the_least_valuable_unprotected_entry() {
+        let directory = temporary_directory("evict");
+        let store = KvStore::with_budget(&directory, 300).unwrap();
+        let first = add(&store, b"first", 10, 32);
+        let second = add(&store, b"second", 100, 32);
+        let third = add(&store, b"third", 200, 32);
+
+        assert!(!first.checkpoint.exists());
+        assert!(second.checkpoint.exists());
+        assert!(third.checkpoint.exists());
+
+        fs::remove_dir_all(directory).unwrap();
+    }
+}