Manage the KV cache disc usage from preferences and stats

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

src/engine/kvstore.rs

@@ -7,7 +7,6 @@ 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;
 
@@ -58,23 +57,19 @@ pub(super) struct KvStore {
 }
 
 impl KvStore {
-    pub(super) fn open(directory: &Path) -> Result<Self, String> {
+    pub(super) fn open(directory: &Path, budget_bytes: u64) -> Result<Self, String> {
         fs::create_dir_all(directory).map_err(|error| {
             format!(
                 "Could not create transient KV cache directory {}: {error}",
                 directory.display()
             )
         })?;
-        Ok(Self {
+        let store = 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;
+            budget_bytes,
+        };
+        // A lowered budget takes effect at once, like ds4's eviction on open.
+        store.evict(None);
         Ok(store)
     }
 
@@ -316,6 +311,7 @@ fn read_u64(file: &mut File) -> Result<u64, String> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::settings::GIB;
 
     fn temporary_directory(name: &str) -> PathBuf {
         let path = std::env::temp_dir().join(format!(
@@ -348,7 +344,7 @@ mod tests {
     #[test]
     fn longest_compatible_prefix_wins_and_hits_survive_reopen() {
         let directory = temporary_directory("prefix");
-        let store = KvStore::open(&directory).unwrap();
+        let store = KvStore::open(&directory, GIB).unwrap();
         add(&store, b"system", 100, 32);
         let expected = add(&store, b"system-user-one", 200, 32);
 
@@ -356,7 +352,7 @@ mod tests {
         assert_eq!(found.checkpoint, expected.checkpoint);
         store.touch(&found).unwrap();
         assert_eq!(
-            KvStore::open(&directory)
+            KvStore::open(&directory, GIB)
                 .unwrap()
                 .find(b"system-user-one-more", 4096)
                 .unwrap()
@@ -371,7 +367,7 @@ mod tests {
     #[test]
     fn budget_evicts_the_least_valuable_unprotected_entry() {
         let directory = temporary_directory("evict");
-        let store = KvStore::with_budget(&directory, 300).unwrap();
+        let store = KvStore::open(&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);

src/engine/metal.rs

@@ -801,6 +801,12 @@ impl Executor {
         self.checkpoint_tag
     }
 
+    /// Marks what the live KV state represents when no checkpoint file is
+    /// written for it, so the next turn can still continue in memory.
+    pub(super) fn note_checkpoint_tag(&mut self, tag: [u8; 32]) {
+        self.checkpoint_tag = tag;
+    }
+
     fn encode_token(&mut self, token: u32) -> Result<(), String> {
         let shape = self.model.shape;
         let map = self.model.main.map_ptr().cast();