1use std::sync::Arc;
16
17use tokio::time::{Duration, interval};
18use tokio_util::sync::CancellationToken;
19
20use crate::embedding_store::EmbeddingStore;
21use crate::error::MemoryError;
22use crate::sqlite_time::parse_sqlite_datetime_to_unix;
23use crate::store::SqliteStore;
24use crate::types::MessageId;
25
26#[derive(Debug, Clone)]
30pub struct EvictionEntry {
31 pub id: MessageId,
33 pub created_at: String,
35 pub last_accessed: Option<String>,
37 pub access_count: u32,
39}
40
41pub trait EvictionPolicy: Send + Sync {
45 fn score(&self, entry: &EvictionEntry) -> f64;
50}
51
52use zeph_config::memory::EvictionConfig;
53
54pub struct EbbinghausPolicy {
69 retention_strength: f64,
70}
71
72impl EbbinghausPolicy {
73 #[must_use]
77 pub fn new(retention_strength: f64) -> Self {
78 Self { retention_strength }
79 }
80}
81
82impl Default for EbbinghausPolicy {
83 fn default() -> Self {
84 Self::new(86_400.0) }
86}
87
88impl EvictionPolicy for EbbinghausPolicy {
89 fn score(&self, entry: &EvictionEntry) -> f64 {
90 let now_secs = unix_now_secs();
91
92 let reference_secs = entry
93 .last_accessed
94 .as_deref()
95 .and_then(parse_sqlite_timestamp_secs)
96 .unwrap_or_else(|| parse_sqlite_timestamp_secs(&entry.created_at).unwrap_or(now_secs));
97
98 #[allow(clippy::cast_precision_loss)]
100 let t = now_secs.saturating_sub(reference_secs) as f64;
101 let n = f64::from(entry.access_count);
102
103 let denominator = (self.retention_strength * (1.0_f64 + n).ln()).max(1.0);
105 (-t / denominator).exp()
106 }
107}
108
109fn unix_now_secs() -> u64 {
110 std::time::SystemTime::now()
111 .duration_since(std::time::UNIX_EPOCH)
112 .map_or(0, |d| d.as_secs())
113}
114
115fn parse_sqlite_timestamp_secs(s: &str) -> Option<u64> {
135 u64::try_from(parse_sqlite_datetime_to_unix(s)?).ok()
136}
137
138#[tracing::instrument(name = "memory.eviction.start_loop", skip_all)]
158pub async fn start_eviction_loop(
159 store: Arc<SqliteStore>,
160 embedding: Option<Arc<EmbeddingStore>>,
161 config: EvictionConfig,
162 policy: Arc<dyn EvictionPolicy + 'static>,
163 cancel: CancellationToken,
164) {
165 if config.max_entries == 0 {
166 tracing::debug!("eviction disabled (max_entries = 0)");
167 return;
168 }
169
170 let mut ticker = interval(Duration::from_secs(config.sweep_interval_secs));
171 ticker.tick().await;
173
174 loop {
175 tokio::select! {
176 () = cancel.cancelled() => {
177 tracing::debug!("eviction loop shutting down");
178 return;
179 }
180 _ = ticker.tick() => {}
181 }
182
183 tracing::debug!(max_entries = config.max_entries, "running eviction sweep");
184
185 match run_eviction_phase1(&store, &*policy, config.max_entries).await {
187 Ok(deleted) => {
188 if deleted > 0 {
189 tracing::info!(deleted, "eviction phase 1: soft-deleted entries");
190 }
191 }
192 Err(e) => {
193 tracing::warn!(error = %e, "eviction phase 1 failed, will retry next sweep");
194 }
195 }
196
197 match run_eviction_phase2(&store, embedding.as_deref()).await {
200 Ok(cleaned) => {
201 if cleaned > 0 {
202 tracing::info!(cleaned, "eviction phase 2: removed Qdrant vectors");
203 }
204 }
205 Err(e) => {
206 tracing::warn!(error = %e, "eviction phase 2 failed, will retry next sweep");
207 }
208 }
209 }
210}
211
212#[cfg_attr(
213 feature = "profiling",
214 tracing::instrument(name = "memory.eviction_phase1", skip_all)
215)]
216async fn run_eviction_phase1(
217 store: &SqliteStore,
218 policy: &dyn EvictionPolicy,
219 max_entries: usize,
220) -> Result<usize, MemoryError> {
221 let candidates = store.get_eviction_candidates().await?;
222 let total = candidates.len();
223
224 if total <= max_entries {
225 return Ok(0);
226 }
227
228 let excess = total - max_entries;
229 let mut scored: Vec<(f64, MessageId)> = candidates
230 .into_iter()
231 .map(|e| (policy.score(&e), e.id))
232 .collect();
233
234 scored.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
236
237 let candidates_to_delete: Vec<MessageId> =
238 scored.into_iter().take(excess).map(|(_, id)| id).collect();
239 let candidate_count = candidates_to_delete.len();
240 let ids_to_delete = store
245 .filter_out_preserved_episode_ids(&candidates_to_delete)
246 .await?;
247 let preserved = candidate_count - ids_to_delete.len();
248 if preserved > 0 {
249 tracing::debug!(
250 preserved,
251 deleted = ids_to_delete.len(),
252 "eviction phase 1: {preserved} candidate(s) preserved by summary anchor"
253 );
254 }
255 store.soft_delete_messages(&ids_to_delete).await?;
256
257 Ok(ids_to_delete.len())
259}
260
261#[cfg_attr(
262 feature = "profiling",
263 tracing::instrument(name = "memory.eviction_phase2", skip_all)
264)]
265async fn run_eviction_phase2(
266 store: &SqliteStore,
267 embedding: Option<&EmbeddingStore>,
268) -> Result<usize, MemoryError> {
269 let ids = store.get_soft_deleted_message_ids().await?;
271 if ids.is_empty() {
272 return Ok(0);
273 }
274
275 if let Some(emb) = embedding {
276 emb.delete_by_message_ids(&ids).await?;
279 } else {
280 tracing::debug!(
281 count = ids.len(),
282 "eviction phase 2: Qdrant disabled, cleaning SQLite bookkeeping only"
283 );
284 }
285
286 store.mark_qdrant_cleaned(&ids).await?;
287 Ok(ids.len())
288}
289
290#[cfg(test)]
293mod tests {
294 use super::*;
295
296 fn ts_ago(seconds_ago: u64) -> String {
300 let ts = unix_now_secs().saturating_sub(seconds_ago);
301 let sec = ts % 60;
303 let min = (ts / 60) % 60;
304 let hour = (ts / 3600) % 24;
305 let mut total_days = ts / 86400;
306 let is_leap =
307 |y: u64| (y.is_multiple_of(4) && !y.is_multiple_of(100)) || y.is_multiple_of(400);
308 let mut year = 1970u64;
309 loop {
310 let days_in_year = if is_leap(year) { 366 } else { 365 };
311 if total_days < days_in_year {
312 break;
313 }
314 total_days -= days_in_year;
315 year += 1;
316 }
317 let month_days = [
318 0u64,
319 31,
320 28 + u64::from(is_leap(year)),
321 31,
322 30,
323 31,
324 30,
325 31,
326 31,
327 30,
328 31,
329 30,
330 31,
331 ];
332 let mut month = 1u64;
333 while month <= 12 {
334 let month_idx = usize::try_from(month).unwrap_or_else(|_| unreachable!());
335 if total_days < month_days[month_idx] {
336 break;
337 }
338 total_days -= month_days[month_idx];
339 month += 1;
340 }
341 let day = total_days + 1;
342 format!("{year:04}-{month:02}-{day:02} {hour:02}:{min:02}:{sec:02}")
343 }
344
345 fn make_entry(access_count: u32, seconds_ago: u64) -> EvictionEntry {
346 let ts = ts_ago(seconds_ago);
347 EvictionEntry {
348 id: MessageId(1),
349 created_at: ts.clone(),
350 last_accessed: Some(ts),
351 access_count,
352 }
353 }
354
355 #[test]
356 fn ebbinghaus_recent_high_access_scores_near_one() {
357 let policy = EbbinghausPolicy::default();
358 let entry = make_entry(10, 1);
360 let score = policy.score(&entry);
361 assert!(
363 score > 0.99,
364 "score should be near 1.0 for recently accessed entry, got {score}"
365 );
366 }
367
368 #[test]
369 fn ebbinghaus_old_zero_access_scores_lower() {
370 let policy = EbbinghausPolicy::default();
371 let old = make_entry(0, 7 * 24 * 3600); let recent = make_entry(0, 60); assert!(
374 policy.score(&old) < policy.score(&recent),
375 "old entry must score lower than recent"
376 );
377 }
378
379 #[test]
380 fn ebbinghaus_high_access_decays_slower() {
381 let policy = EbbinghausPolicy::default();
382 let low = make_entry(1, 3600); let high = make_entry(20, 3600); assert!(
385 policy.score(&high) > policy.score(&low),
386 "high access count should yield higher score"
387 );
388 }
389
390 #[test]
391 fn ebbinghaus_never_accessed_uses_created_at_as_reference() {
392 let policy = EbbinghausPolicy::default();
393 let old_with_no_last_accessed = EvictionEntry {
396 id: MessageId(2),
397 created_at: ts_ago(7 * 24 * 3600),
398 last_accessed: None,
399 access_count: 0,
400 };
401 let old_with_same_last_accessed = make_entry(0, 7 * 24 * 3600);
402 let score_no_access = policy.score(&old_with_no_last_accessed);
403 let score_same = policy.score(&old_with_same_last_accessed);
404 let diff = (score_no_access - score_same).abs();
406 assert!(diff < 1e-6, "scores should match; diff = {diff}");
407 }
408
409 #[test]
410 fn eviction_config_default_is_disabled() {
411 let config = EvictionConfig::default();
412 assert_eq!(
413 config.max_entries, 0,
414 "eviction must be disabled by default"
415 );
416 }
417
418 #[tokio::test]
421 async fn test_eviction_preserves_summary_anchored_messages() {
422 use crate::store::SqliteStore;
423
424 let store = SqliteStore::new(":memory:").await.unwrap();
425 let cid = store.create_conversation().await.unwrap();
426
427 let ids: Vec<_> = (0..6)
429 .map(|_| async { store.save_message(cid, "user", "msg").await.unwrap() })
430 .collect();
431 let mut msg_ids = Vec::new();
432 for f in ids {
433 msg_ids.push(f.await);
434 }
435
436 store
438 .save_summary(cid, "summary", Some(msg_ids[0]), Some(msg_ids[2]), 30)
439 .await
440 .unwrap();
441
442 let policy = EbbinghausPolicy::default();
443 let deleted = run_eviction_phase1(&store, &policy, 3).await.unwrap();
445
446 assert!(
449 deleted <= 3,
450 "at most 3 messages can be deleted, got {deleted}"
451 );
452
453 for &anchored in &msg_ids[0..=2] {
454 let is_deleted: Option<String> =
455 sqlx::query_scalar("SELECT deleted_at FROM messages WHERE id = ?")
456 .bind(anchored)
457 .fetch_one(store.pool())
458 .await
459 .unwrap();
460 assert!(
461 is_deleted.is_none(),
462 "summary-anchored message {anchored:?} must not be soft-deleted"
463 );
464 }
465 }
466
467 #[test]
468 fn parse_sqlite_timestamp_known_value() {
469 let ts = parse_sqlite_timestamp_secs("2024-01-01 00:00:00").unwrap();
471 assert_eq!(
474 ts, 1_704_067_200,
475 "2024-01-01 must parse to known timestamp"
476 );
477 }
478
479 #[test]
480 fn parse_sqlite_timestamp_pre_1970_returns_none() {
481 assert_eq!(parse_sqlite_timestamp_secs("1969-12-31 23:59:59"), None);
485 assert_eq!(parse_sqlite_timestamp_secs("1900-01-01 00:00:00"), None);
486 }
487
488 async fn setup_embedding_store() -> (EmbeddingStore, crate::store::SqliteStore) {
493 let sqlite = crate::store::SqliteStore::new(":memory:").await.unwrap();
494 let pool = sqlite.pool().clone();
495 let mem_store = Box::new(crate::in_memory_store::InMemoryVectorStore::new());
496 let emb = EmbeddingStore::with_store(mem_store, pool);
497 emb.ensure_collection(4).await.unwrap();
498 (emb, sqlite)
499 }
500
501 async fn seed_soft_deleted(
503 store: &crate::store::SqliteStore,
504 emb: &EmbeddingStore,
505 ) -> (MessageId, String) {
506 let cid = store.create_conversation().await.unwrap();
507 let msg_id = store.save_message(cid, "user", "hello").await.unwrap();
508
509 let point_id = emb
510 .store(
511 msg_id,
512 cid,
513 "user",
514 vec![1.0, 0.0, 0.0, 0.0],
515 crate::embedding_store::MessageKind::Regular,
516 "test",
517 0,
518 )
519 .await
520 .unwrap();
521
522 store.soft_delete_messages(&[msg_id]).await.unwrap();
524
525 (msg_id, point_id)
526 }
527
528 #[tokio::test]
530 async fn eviction_phase2_calls_delete_before_mark_clean() {
531 let (emb, store) = setup_embedding_store().await;
532 let (msg_id, _point_id) = seed_soft_deleted(&store, &emb).await;
533
534 let cleaned = run_eviction_phase2(&store, Some(&emb)).await.unwrap();
536 assert_eq!(cleaned, 1, "one message should be cleaned");
537
538 let remaining = store.get_soft_deleted_message_ids().await.unwrap();
540 assert!(
541 !remaining.contains(&msg_id),
542 "message must not remain in soft-deleted list after phase 2"
543 );
544 }
545
546 #[tokio::test]
548 async fn eviction_phase2_skips_delete_when_no_embedding_store() {
549 let (_, store) = setup_embedding_store().await;
550 let cid = store.create_conversation().await.unwrap();
551 let msg_id = store.save_message(cid, "user", "hello").await.unwrap();
552 store.soft_delete_messages(&[msg_id]).await.unwrap();
553
554 let cleaned = run_eviction_phase2(&store, None).await.unwrap();
556 assert_eq!(
557 cleaned, 1,
558 "row must be cleaned even without embedding store"
559 );
560
561 let remaining = store.get_soft_deleted_message_ids().await.unwrap();
562 assert!(
563 !remaining.contains(&msg_id),
564 "message must not remain in soft-deleted list"
565 );
566 }
567
568 #[tokio::test]
570 async fn eviction_phase2_empty_returns_zero() {
571 let (emb, store) = setup_embedding_store().await;
572 let cleaned = run_eviction_phase2(&store, Some(&emb)).await.unwrap();
573 assert_eq!(cleaned, 0, "no soft-deleted messages → 0 cleaned");
574 }
575}