| 1 | //! Persistent memory snapshots for capacity controller interventions. |
| 2 | |
| 3 | use std::fs::{self, OpenOptions}; |
| 4 | use std::io::{BufRead, BufReader, Write}; |
| 5 | use std::path::{Path, PathBuf}; |
| 6 | use std::time::SystemTime; |
| 7 | |
| 8 | use anyhow::{Context, Result, anyhow}; |
| 9 | use chrono::Utc; |
| 10 | use serde::{Deserialize, Serialize}; |
| 11 | |
| 12 | /// Canonical compact state persisted by interventions. |
| 13 | #[derive(Debug, Clone, Serialize, Deserialize, Default)] |
| 14 | pub struct CanonicalState { |
| 15 | pub goal: String, |
| 16 | pub constraints: Vec<String>, |
| 17 | pub confirmed_facts: Vec<String>, |
| 18 | pub open_loops: Vec<String>, |
| 19 | pub pending_actions: Vec<String>, |
| 20 | pub critical_refs: Vec<String>, |
| 21 | } |
| 22 | |
| 23 | /// Replay verification metadata. |
| 24 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 25 | pub struct ReplayInfo { |
| 26 | pub tool_id: String, |
| 27 | pub tool_name: String, |
| 28 | pub pass: bool, |
| 29 | pub diff_summary: String, |
| 30 | } |
| 31 | |
| 32 | /// JSONL record written for each intervention. |
| 33 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 34 | pub struct CapacityMemoryRecord { |
| 35 | pub id: String, |
| 36 | pub ts: String, |
| 37 | pub turn_index: u64, |
| 38 | pub action_trigger: String, |
| 39 | pub h_hat: f64, |
| 40 | pub c_hat: f64, |
| 41 | pub slack: f64, |
| 42 | pub risk_band: String, |
| 43 | pub canonical_state: CanonicalState, |
| 44 | pub source_message_ids: Vec<String>, |
| 45 | #[serde(skip_serializing_if = "Option::is_none")] |
| 46 | pub replay_info: Option<ReplayInfo>, |
| 47 | } |
| 48 | |
| 49 | fn capacity_memory_dirs() -> Vec<PathBuf> { |
| 50 | if let Ok(raw) = std::env::var("DEEPSEEK_CAPACITY_MEMORY_DIR") { |
| 51 | let trimmed = raw.trim(); |
| 52 | if !trimmed.is_empty() { |
| 53 | return vec![PathBuf::from(shellexpand::tilde(trimmed).as_ref())]; |
| 54 | } |
| 55 | } |
| 56 | |
| 57 | let mut dirs = Vec::new(); |
| 58 | if let Some(home) = dirs::home_dir() { |
| 59 | dirs.push(home.join(".deepseek").join("memory")); |
| 60 | } |
| 61 | |
| 62 | let cwd = std::env::current_dir() |
| 63 | .unwrap_or_else(|_| PathBuf::from(".")) |
| 64 | .join(".deepseek") |
| 65 | .join("memory"); |
| 66 | dirs.push(cwd); |
| 67 | |
| 68 | dirs.dedup(); |
| 69 | dirs |
| 70 | } |
| 71 | |
| 72 | pub fn append_capacity_record(session_id: &str, record: &CapacityMemoryRecord) -> Result<PathBuf> { |
| 73 | let candidates = candidate_session_memory_paths(session_id); |
| 74 | append_capacity_record_to_candidates(&candidates, record) |
| 75 | } |
| 76 | |
| 77 | pub fn append_capacity_record_to_path(path: &Path, record: &CapacityMemoryRecord) -> Result<()> { |
| 78 | if let Some(parent) = path.parent() { |
| 79 | fs::create_dir_all(parent) |
| 80 | .with_context(|| format!("Failed to create memory directory {}", parent.display()))?; |
| 81 | } |
| 82 | let mut file = OpenOptions::new() |
| 83 | .create(true) |
| 84 | .append(true) |
| 85 | .open(path) |
| 86 | .with_context(|| format!("Failed to open memory log {}", path.display()))?; |
| 87 | let line = |
| 88 | serde_json::to_string(record).context("Failed to serialize capacity memory record")?; |
| 89 | writeln!(file, "{line}") |
| 90 | .with_context(|| format!("Failed to write memory record {}", path.display()))?; |
| 91 | Ok(()) |
| 92 | } |
| 93 | |
| 94 | pub fn load_last_k_capacity_records( |
| 95 | session_id: &str, |
| 96 | k: usize, |
| 97 | ) -> Result<Vec<CapacityMemoryRecord>> { |
| 98 | let candidates = candidate_session_memory_paths(session_id); |
| 99 | load_last_k_capacity_records_from_candidates(&candidates, k) |
| 100 | } |
| 101 | |
| 102 | pub fn load_last_k_capacity_records_from_path( |
| 103 | path: &Path, |
| 104 | k: usize, |
| 105 | ) -> Result<Vec<CapacityMemoryRecord>> { |
| 106 | if k == 0 || !path.exists() { |
| 107 | return Ok(Vec::new()); |
| 108 | } |
| 109 | |
| 110 | let file = OpenOptions::new() |
| 111 | .read(true) |
| 112 | .open(path) |
| 113 | .with_context(|| format!("Failed to open memory log {}", path.display()))?; |
| 114 | let reader = BufReader::new(file); |
| 115 | let mut records = Vec::new(); |
| 116 | |
| 117 | for line in reader.lines() { |
| 118 | let line = line.with_context(|| format!("Failed reading {}", path.display()))?; |
| 119 | if line.trim().is_empty() { |
| 120 | continue; |
| 121 | } |
| 122 | if let Ok(record) = serde_json::from_str::<CapacityMemoryRecord>(&line) { |
| 123 | records.push(record); |
| 124 | } |
| 125 | } |
| 126 | |
| 127 | if records.len() > k { |
| 128 | Ok(records.split_off(records.len() - k)) |
| 129 | } else { |
| 130 | Ok(records) |
| 131 | } |
| 132 | } |
| 133 | |
| 134 | fn candidate_session_memory_paths(session_id: &str) -> Vec<PathBuf> { |
| 135 | capacity_memory_dirs() |
| 136 | .into_iter() |
| 137 | .map(|dir| dir.join(format!("{session_id}.jsonl"))) |
| 138 | .collect() |
| 139 | } |
| 140 | |
| 141 | fn append_capacity_record_to_candidates( |
| 142 | paths: &[PathBuf], |
| 143 | record: &CapacityMemoryRecord, |
| 144 | ) -> Result<PathBuf> { |
| 145 | let mut last_err: Option<anyhow::Error> = None; |
| 146 | for path in paths { |
| 147 | match append_capacity_record_to_path(path, record) { |
| 148 | Ok(()) => return Ok(path.clone()), |
| 149 | Err(err) => last_err = Some(err), |
| 150 | } |
| 151 | } |
| 152 | |
| 153 | Err(last_err.unwrap_or_else(|| anyhow!("No capacity memory path candidates available"))) |
| 154 | } |
| 155 | |
| 156 | fn load_last_k_capacity_records_from_candidates( |
| 157 | paths: &[PathBuf], |
| 158 | k: usize, |
| 159 | ) -> Result<Vec<CapacityMemoryRecord>> { |
| 160 | if k == 0 { |
| 161 | return Ok(Vec::new()); |
| 162 | } |
| 163 | |
| 164 | let mut newest: Option<(SystemTime, Vec<CapacityMemoryRecord>)> = None; |
| 165 | let mut last_err: Option<anyhow::Error> = None; |
| 166 | |
| 167 | for path in paths { |
| 168 | if !path.exists() { |
| 169 | continue; |
| 170 | } |
| 171 | |
| 172 | match load_last_k_capacity_records_from_path(path, k) { |
| 173 | Ok(records) => { |
| 174 | if records.is_empty() { |
| 175 | continue; |
| 176 | } |
| 177 | let modified = fs::metadata(path) |
| 178 | .and_then(|meta| meta.modified()) |
| 179 | .unwrap_or(SystemTime::UNIX_EPOCH); |
| 180 | let should_replace = newest |
| 181 | .as_ref() |
| 182 | .map(|(current, _)| modified >= *current) |
| 183 | .unwrap_or(true); |
| 184 | if should_replace { |
| 185 | newest = Some((modified, records)); |
| 186 | } |
| 187 | } |
| 188 | Err(err) => last_err = Some(err), |
| 189 | } |
| 190 | } |
| 191 | |
| 192 | if let Some((_, records)) = newest { |
| 193 | return Ok(records); |
| 194 | } |
| 195 | if let Some(err) = last_err { |
| 196 | return Err(err); |
| 197 | } |
| 198 | Ok(Vec::new()) |
| 199 | } |
| 200 | |
| 201 | #[must_use] |
| 202 | pub fn new_record_id() -> String { |
| 203 | format!("cap_{}", &uuid::Uuid::new_v4().to_string()[..8]) |
| 204 | } |
| 205 | |
| 206 | #[must_use] |
| 207 | pub fn now_rfc3339() -> String { |
| 208 | Utc::now().to_rfc3339() |
| 209 | } |
| 210 | |
| 211 | #[cfg(test)] |
| 212 | mod tests { |
| 213 | use super::*; |
| 214 | use tempfile::tempdir; |
| 215 | |
| 216 | #[test] |
| 217 | fn memory_jsonl_round_trip() { |
| 218 | let tmp = tempdir().expect("tempdir"); |
| 219 | let path = tmp.path().join("session.jsonl"); |
| 220 | |
| 221 | let record = CapacityMemoryRecord { |
| 222 | id: "cap_1".to_string(), |
| 223 | ts: now_rfc3339(), |
| 224 | turn_index: 2, |
| 225 | action_trigger: "targeted_context_refresh".to_string(), |
| 226 | h_hat: 1.2, |
| 227 | c_hat: 3.8, |
| 228 | slack: 2.6, |
| 229 | risk_band: "medium".to_string(), |
| 230 | canonical_state: CanonicalState { |
| 231 | goal: "Ship feature".to_string(), |
| 232 | ..CanonicalState::default() |
| 233 | }, |
| 234 | source_message_ids: vec!["m1".to_string()], |
| 235 | replay_info: None, |
| 236 | }; |
| 237 | |
| 238 | append_capacity_record_to_path(&path, &record).expect("append"); |
| 239 | let records = load_last_k_capacity_records_from_path(&path, 1).expect("load"); |
| 240 | assert_eq!(records.len(), 1); |
| 241 | assert_eq!(records[0].canonical_state.goal, "Ship feature"); |
| 242 | } |
| 243 | |
| 244 | #[test] |
| 245 | fn append_falls_back_to_next_candidate_path() { |
| 246 | let tmp = tempdir().expect("tempdir"); |
| 247 | let blocked_root = tmp.path().join("blocked"); |
| 248 | fs::write(&blocked_root, "file").expect("create blocking file"); |
| 249 | let blocked_path = blocked_root.join("session.jsonl"); |
| 250 | let fallback_path = tmp.path().join("fallback").join("session.jsonl"); |
| 251 | |
| 252 | let record = CapacityMemoryRecord { |
| 253 | id: "cap_fallback".to_string(), |
| 254 | ts: now_rfc3339(), |
| 255 | turn_index: 1, |
| 256 | action_trigger: "targeted_context_refresh".to_string(), |
| 257 | h_hat: 1.0, |
| 258 | c_hat: 3.8, |
| 259 | slack: 2.8, |
| 260 | risk_band: "medium".to_string(), |
| 261 | canonical_state: CanonicalState::default(), |
| 262 | source_message_ids: vec!["m1".to_string()], |
| 263 | replay_info: None, |
| 264 | }; |
| 265 | |
| 266 | let chosen = append_capacity_record_to_candidates( |
| 267 | &[blocked_path.clone(), fallback_path.clone()], |
| 268 | &record, |
| 269 | ) |
| 270 | .expect("append with fallback"); |
| 271 | assert_eq!(chosen, fallback_path); |
| 272 | assert!(chosen.exists()); |
| 273 | } |
| 274 | |
| 275 | #[test] |
| 276 | fn load_prefers_newest_candidate_records() { |
| 277 | let tmp = tempdir().expect("tempdir"); |
| 278 | let older = tmp.path().join("older.jsonl"); |
| 279 | let newer = tmp.path().join("newer.jsonl"); |
| 280 | |
| 281 | let old_record = CapacityMemoryRecord { |
| 282 | id: "cap_old".to_string(), |
| 283 | ts: now_rfc3339(), |
| 284 | turn_index: 1, |
| 285 | action_trigger: "targeted_context_refresh".to_string(), |
| 286 | h_hat: 1.0, |
| 287 | c_hat: 3.8, |
| 288 | slack: 2.8, |
| 289 | risk_band: "medium".to_string(), |
| 290 | canonical_state: CanonicalState { |
| 291 | goal: "old".to_string(), |
| 292 | ..CanonicalState::default() |
| 293 | }, |
| 294 | source_message_ids: vec!["m1".to_string()], |
| 295 | replay_info: None, |
| 296 | }; |
| 297 | let new_record = CapacityMemoryRecord { |
| 298 | id: "cap_new".to_string(), |
| 299 | ts: now_rfc3339(), |
| 300 | turn_index: 2, |
| 301 | action_trigger: "verify_and_replan".to_string(), |
| 302 | h_hat: 1.4, |
| 303 | c_hat: 3.8, |
| 304 | slack: 2.4, |
| 305 | risk_band: "high".to_string(), |
| 306 | canonical_state: CanonicalState { |
| 307 | goal: "new".to_string(), |
| 308 | ..CanonicalState::default() |
| 309 | }, |
| 310 | source_message_ids: vec!["m2".to_string()], |
| 311 | replay_info: None, |
| 312 | }; |
| 313 | |
| 314 | append_capacity_record_to_path(&older, &old_record).expect("write older"); |
| 315 | std::thread::sleep(std::time::Duration::from_millis(10)); |
| 316 | append_capacity_record_to_path(&newer, &new_record).expect("write newer"); |
| 317 | |
| 318 | let records = load_last_k_capacity_records_from_candidates(&[older, newer], 1) |
| 319 | .expect("load newest records"); |
| 320 | assert_eq!(records.len(), 1); |
| 321 | assert_eq!(records[0].canonical_state.goal, "new"); |
| 322 | } |
| 323 | } |
| 324 |