//! Background persistence for tool state. //! //! [`ResourcesPersistence`] persists `Resources` state (the new architecture). //! Old `ToolStatePersistence` and `PersistenceLayer` have been deleted. use std::path::{Path, PathBuf}; use std::time::Duration; use crate::types::resources::Resources; /// Background persistence for `Resources` state/params. /// /// Same pattern as `ToolStatePersistence` — debounced background writes with /// atomic rename. Takes a `serde_json::Value` from `Resources::serialize()` /// and writes it to disk. On load, parses the JSON and feeds it to /// `Resources::load_from()`. /// /// This replaces the old `ToolStatePersistence` pipeline for the new /// architecture. During migration both coexist; once all tools are migrated, /// `ToolStatePersistence` will be deleted. pub struct ResourcesPersistence { /// Path to the JSON file where Resources state is persisted state_path: PathBuf, /// Channel to send serialized state to the background writer tx: tokio::sync::mpsc::UnboundedSender, } enum ResourcesPersistenceCommand { /// Write this serialized Resources value to disk Save(serde_json::Value), /// Flush pending writes and notify when done Flush(tokio::sync::oneshot::Sender<()>), } impl ResourcesPersistence { /// Construct a noop persistence handle for tests. No background task. pub fn noop() -> Self { let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); Self { state_path: PathBuf::from("/dev/null"), tx, } } /// Create a new persistence handle and spawn the background writer task. pub fn new(state_path: PathBuf) -> Self { let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); let writer_path = state_path.clone(); tokio::spawn(async move { Self::writer_loop(rx, writer_path).await; }); Self { state_path, tx } } /// Load existing Resources state from disk, if the file exists. /// /// Reads the JSON, parses it into the nested `HashMap>` /// shape that `Resources::load_from()` expects, and applies it to the given resources. /// /// Returns `true` if state was loaded, `false` if no file or parse error. pub fn load(&self, resources: &mut Resources) -> bool { let json = match std::fs::read_to_string(&self.state_path) { Ok(s) => s, Err(_) => return false, }; let top: serde_json::Value = match serde_json::from_str(&json) { Ok(v) => v, Err(e) => { tracing::warn!( "Failed to parse resources state from {:?}: {}", self.state_path, e ); return false; } }; let data = match Self::value_to_nested_map(top) { Some(m) => m, None => { tracing::warn!( "Resources state file {:?} has unexpected shape", self.state_path ); return false; } }; resources.load_from(data); true } /// Save the current Resources state (non-blocking). /// Sends a serialized snapshot to the background writer. pub fn save(&self, resources: &Resources) { let snapshot = resources.serialize(); let _ = self.tx.send(ResourcesPersistenceCommand::Save(snapshot)); } /// Path to the persisted state file. pub fn state_path(&self) -> &std::path::Path { &self.state_path } /// Flush pending writes. Call on graceful shutdown. pub async fn flush(&self) { let (done_tx, done_rx) = tokio::sync::oneshot::channel(); let _ = self.tx.send(ResourcesPersistenceCommand::Flush(done_tx)); let _ = done_rx.await; } /// Convert the `serde_json::Value` (from serialize()) into the nested /// HashMap structure that `load_from()` expects. fn value_to_nested_map( val: serde_json::Value, ) -> Option< std::collections::HashMap>, > { let top = val.as_object()?; let mut result = std::collections::HashMap::new(); for (cat_key, cat_val) in top { let inner_obj = cat_val.as_object()?; let mut inner = std::collections::HashMap::new(); for (k, v) in inner_obj { inner.insert(k.clone(), v.clone()); } result.insert(cat_key.clone(), inner); } Some(result) } async fn writer_loop( mut rx: tokio::sync::mpsc::UnboundedReceiver, state_path: PathBuf, ) { let mut pending: Option = None; let mut debounce = tokio::time::interval(Duration::from_millis(500)); debounce.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); loop { tokio::select! { cmd = rx.recv() => { match cmd { Some(ResourcesPersistenceCommand::Save(snapshot)) => { pending = Some(snapshot); } Some(ResourcesPersistenceCommand::Flush(done)) => { if let Some(snapshot) = pending.take() { Self::write_json(&state_path, &snapshot).await; } let _ = done.send(()); } None => { if let Some(snapshot) = pending.take() { Self::write_json(&state_path, &snapshot).await; } break; } } } _ = debounce.tick() => { if let Some(snapshot) = pending.take() { Self::write_json(&state_path, &snapshot).await; } } } } } async fn write_json(path: &Path, value: &serde_json::Value) { match serde_json::to_string_pretty(value) { Ok(json) => { let tmp_path = path.with_extension("json.tmp"); if let Err(e) = tokio::fs::write(&tmp_path, json.as_bytes()).await { tracing::warn!("Failed to write resources state to {:?}: {}", tmp_path, e); return; } // Guard: if a previous bug left a directory at `path`, remove it // so the atomic rename can succeed. if path.is_dir() { tracing::warn!( "Resources state path {:?} is a directory — removing before write", path ); let _ = tokio::fs::remove_dir_all(path).await; } if let Err(e) = tokio::fs::rename(&tmp_path, path).await { tracing::warn!( "Failed to rename resources state {:?} -> {:?}: {}", tmp_path, path, e ); } } Err(e) => { tracing::warn!("Failed to serialize resources state: {}", e); } } } } // Old `PersistenceLayer` / `PersistenceRunner` deleted. // ToolState persistence replaced by Resources persistence via `ResourcesPersistence`. #[cfg(test)] mod tests { use super::*; // ResourcesPersistence tests // ----------------------------------------------------------------------- use crate::types::resources::{Resources, State, WebCitationCounter}; #[tokio::test] async fn resources_save_and_load_roundtrip() { let dir = tempfile::tempdir().unwrap(); let state_path = dir.path().join("resources_state.json"); let persistence = ResourcesPersistence::new(state_path); // Build resources with registered state types let mut resources = Resources::new(); resources.register_state::(); // Populate WebCitationCounter { let counter = resources.get_or_default::>(); counter.counter = 7; } // Save and flush persistence.save(&resources); persistence.flush().await; // Load into fresh resources (with same registrations) let mut restored = Resources::new(); restored.register_state::(); assert!(persistence.load(&mut restored)); // Verify WebCitationCounter roundtripped let counter = restored.get::>().unwrap(); assert_eq!(counter.counter, 7); } #[tokio::test] async fn resources_load_returns_false_when_no_file() { let dir = tempfile::tempdir().unwrap(); let state_path = dir.path().join("nonexistent.json"); let persistence = ResourcesPersistence::new(state_path); let mut resources = Resources::new(); assert!(!persistence.load(&mut resources)); } #[tokio::test] async fn resources_load_returns_false_on_corrupt_json() { let dir = tempfile::tempdir().unwrap(); let state_path = dir.path().join("resources_state.json"); std::fs::write(&state_path, "{ this is not valid json }").unwrap(); let persistence = ResourcesPersistence::new(state_path); let mut resources = Resources::new(); assert!(!persistence.load(&mut resources)); } /// Atomic-rename guarantee: a concurrent reader hammering the path while /// the writer streams 200 snapshots must never observe torn JSON. #[tokio::test] async fn writer_atomic_rename_never_exposes_partial_json() { use std::sync::atomic::{AtomicBool, Ordering}; let dir = tempfile::tempdir().unwrap(); let state_path = dir.path().join("resources_state.json"); let persistence = ResourcesPersistence::new(state_path.clone()); let mut resources = Resources::new(); resources.register_state::(); let done = std::sync::Arc::new(AtomicBool::new(false)); let reader_done = done.clone(); let reader_path = state_path.clone(); let reader = tokio::spawn(async move { while !reader_done.load(Ordering::Relaxed) { if let Ok(s) = tokio::fs::read_to_string(&reader_path).await { assert!( serde_json::from_str::(&s).is_ok(), "reader observed a torn/partial write (atomic-rename violated): {s:?}" ); } tokio::task::yield_now().await; } }); for i in 0..200u64 { { let counter = resources.get_or_default::>(); counter.counter = i as u32; } persistence.save(&resources); persistence.flush().await; } done.store(true, Ordering::Relaxed); reader.await.unwrap(); // Final state is intact and reflects the last write. let content = std::fs::read_to_string(&state_path).unwrap(); let parsed: serde_json::Value = serde_json::from_str(&content).unwrap(); assert!(parsed["state"]["grok_build.WebCitation"].is_object()); } #[tokio::test] async fn resources_flush_writes_pending() { let dir = tempfile::tempdir().unwrap(); let state_path = dir.path().join("resources_state.json"); let persistence = ResourcesPersistence::new(state_path.clone()); let mut resources = Resources::new(); resources.register_state::(); { let counter = resources.get_or_default::>(); counter.counter = 42; } persistence.save(&resources); persistence.flush().await; // File should exist with correct structure assert!(state_path.exists()); let content = std::fs::read_to_string(&state_path).unwrap(); let parsed: serde_json::Value = serde_json::from_str(&content).unwrap(); // Should have "state" category with "grok_build.WebCitation" key assert!(parsed["state"]["grok_build.WebCitation"].is_object()); } }