1use zeph_db::{ActiveDialect, query, query_as, query_scalar, sql};
5
6use super::SqliteStore;
7use crate::error::MemoryError;
8
9#[derive(Debug, Clone, sqlx::FromRow)]
11pub struct MemoryTreeRow {
12 pub id: i64,
13 pub level: i64,
14 pub parent_id: Option<i64>,
15 pub content: String,
16 pub source_ids: String,
17 pub token_count: i64,
18 pub consolidated_at: Option<String>,
19 pub created_at: String,
20}
21
22impl SqliteStore {
23 pub async fn insert_tree_leaf(
31 &self,
32 content: &str,
33 token_count: i64,
34 ) -> Result<i64, MemoryError> {
35 let (id,): (i64,) = query_as(sql!(
36 "INSERT INTO memory_tree (level, content, token_count)
37 VALUES (0, ?, ?)
38 RETURNING id"
39 ))
40 .bind(content)
41 .bind(token_count)
42 .fetch_one(self.pool())
43 .await?;
44
45 Ok(id)
46 }
47
48 pub async fn insert_tree_node(
56 &self,
57 level: i64,
58 parent_id: Option<i64>,
59 content: &str,
60 source_ids: &str,
61 token_count: i64,
62 ) -> Result<i64, MemoryError> {
63 let now = <ActiveDialect as zeph_db::dialect::Dialect>::NOW;
64 let raw = format!(
65 "INSERT INTO memory_tree
66 (level, parent_id, content, source_ids, token_count, consolidated_at)
67 VALUES (?, ?, ?, ?, ?, {now})
68 RETURNING id"
69 );
70 let query_sql = zeph_db::rewrite_placeholders(&raw);
71 let (id,): (i64,) = query_as(sqlx::AssertSqlSafe(query_sql))
72 .bind(level)
73 .bind(parent_id)
74 .bind(content)
75 .bind(source_ids)
76 .bind(token_count)
77 .fetch_one(self.pool())
78 .await?;
79
80 Ok(id)
81 }
82
83 pub async fn load_tree_leaves_unconsolidated(
89 &self,
90 limit: usize,
91 ) -> Result<Vec<MemoryTreeRow>, MemoryError> {
92 let consolidated_at_sel =
97 <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("consolidated_at");
98 let created_at_sel =
99 <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("created_at");
100 let raw = format!(
115 "SELECT id, level, parent_id, content, source_ids, token_count,
116 {consolidated_at_sel} AS consolidated_at, {created_at_sel} AS created_at
117 FROM (
118 SELECT id, level, parent_id, content, source_ids, token_count,
119 consolidated_at, created_at
120 FROM memory_tree
121 WHERE level = 0 AND parent_id IS NULL
122 ORDER BY COALESCE(last_attempted_at, created_at) ASC,
123 consolidation_attempts ASC,
124 id ASC
125 LIMIT ?
126 ) AS batch
127 ORDER BY batch.created_at ASC, batch.id ASC"
128 );
129 let query_sql = zeph_db::rewrite_placeholders(&raw);
130 let rows: Vec<MemoryTreeRow> = query_as(sqlx::AssertSqlSafe(query_sql))
131 .bind(i64::try_from(limit).unwrap_or(i64::MAX))
132 .fetch_all(self.pool())
133 .await?;
134
135 Ok(rows)
136 }
137
138 pub async fn load_tree_level(
144 &self,
145 level: i64,
146 limit: usize,
147 ) -> Result<Vec<MemoryTreeRow>, MemoryError> {
148 let consolidated_at_sel =
151 <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("consolidated_at");
152 let created_at_sel =
153 <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("created_at");
154 let raw = format!(
158 "SELECT id, level, parent_id, content, source_ids, token_count,
159 {consolidated_at_sel} AS consolidated_at, {created_at_sel} AS created_at
160 FROM (
161 SELECT id, level, parent_id, content, source_ids, token_count,
162 consolidated_at, created_at
163 FROM memory_tree
164 WHERE level = ? AND parent_id IS NULL
165 ORDER BY COALESCE(last_attempted_at, created_at) ASC,
166 consolidation_attempts ASC,
167 id ASC
168 LIMIT ?
169 ) AS batch
170 ORDER BY batch.created_at ASC, batch.id ASC"
171 );
172 let query_sql = zeph_db::rewrite_placeholders(&raw);
173 let rows: Vec<MemoryTreeRow> = query_as(sqlx::AssertSqlSafe(query_sql))
174 .bind(level)
175 .bind(i64::try_from(limit).unwrap_or(i64::MAX))
176 .fetch_all(self.pool())
177 .await?;
178
179 Ok(rows)
180 }
181
182 pub async fn traverse_tree_up(
190 &self,
191 leaf_id: i64,
192 max_level: i64,
193 ) -> Result<Vec<MemoryTreeRow>, MemoryError> {
194 let mut result = Vec::new();
196 let mut current_id = leaf_id;
197
198 let consolidated_at_sel =
201 <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("consolidated_at");
202 let created_at_sel =
203 <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("created_at");
204 let raw = format!(
205 "SELECT id, level, parent_id, content, source_ids, token_count,
206 {consolidated_at_sel} AS consolidated_at, {created_at_sel} AS created_at
207 FROM memory_tree
208 WHERE id = ?"
209 );
210 let query_sql = zeph_db::rewrite_placeholders(&raw);
211
212 for _ in 0..=max_level {
213 let row: Option<MemoryTreeRow> = query_as(sqlx::AssertSqlSafe(query_sql.clone()))
214 .bind(current_id)
215 .fetch_optional(self.pool())
216 .await?;
217
218 match row {
219 None => break,
220 Some(r) => {
221 let next_id = r.parent_id;
222 result.push(r);
223 match next_id {
224 None => break,
225 Some(p) => current_id = p,
226 }
227 }
228 }
229 }
230
231 Ok(result)
232 }
233
234 pub async fn mark_nodes_consolidated(
244 &self,
245 child_ids: &[i64],
246 parent_id: i64,
247 ) -> Result<(), MemoryError> {
248 if child_ids.is_empty() {
249 return Ok(());
250 }
251
252 let mut tx = self.pool().begin().await?;
253
254 let now = <ActiveDialect as zeph_db::dialect::Dialect>::NOW;
255 let raw = format!(
256 "UPDATE memory_tree
257 SET parent_id = ?, consolidated_at = {now}
258 WHERE id = ? AND parent_id IS NULL"
259 );
260 let query_sql = zeph_db::rewrite_placeholders(&raw);
261 for &child_id in child_ids {
262 query(sqlx::AssertSqlSafe(query_sql.as_str()))
263 .bind(parent_id)
264 .bind(child_id)
265 .execute(&mut *tx)
266 .await?;
267 }
268
269 tx.commit().await?;
270 Ok(())
271 }
272
273 pub async fn bump_consolidation_attempts(&self, node_ids: &[i64]) -> Result<(), MemoryError> {
289 if node_ids.is_empty() {
290 return Ok(());
291 }
292
293 let mut tx = self.pool().begin().await?;
294
295 let now = <ActiveDialect as zeph_db::dialect::Dialect>::NOW;
296 let raw = format!(
297 "UPDATE memory_tree
298 SET consolidation_attempts = consolidation_attempts + 1,
299 last_attempted_at = {now}
300 WHERE id = ?"
301 );
302 let query_sql = zeph_db::rewrite_placeholders(&raw);
303 for &node_id in node_ids {
304 query(sqlx::AssertSqlSafe(query_sql.as_str()))
305 .bind(node_id)
306 .execute(&mut *tx)
307 .await?;
308 }
309
310 tx.commit().await?;
311 Ok(())
312 }
313
314 #[cfg_attr(
323 feature = "profiling",
324 tracing::instrument(name = "memory.consolidate", skip_all)
325 )]
326 pub async fn consolidate_cluster(
327 &self,
328 level: i64,
329 summary: &str,
330 source_ids_json: &str,
331 token_count: i64,
332 child_ids: &[i64],
333 ) -> Result<i64, MemoryError> {
334 if child_ids.is_empty() {
335 return Err(MemoryError::InvalidInput(
336 "child_ids must not be empty".into(),
337 ));
338 }
339
340 let mut tx = self.pool().begin().await?;
341
342 let now = <ActiveDialect as zeph_db::dialect::Dialect>::NOW;
343 let insert_raw = format!(
344 "INSERT INTO memory_tree
345 (level, content, source_ids, token_count, consolidated_at)
346 VALUES (?, ?, ?, ?, {now})
347 RETURNING id"
348 );
349 let insert_sql = zeph_db::rewrite_placeholders(&insert_raw);
350 let (parent_id,): (i64,) = zeph_db::query_as(sqlx::AssertSqlSafe(insert_sql))
351 .bind(level)
352 .bind(summary)
353 .bind(source_ids_json)
354 .bind(token_count)
355 .fetch_one(&mut *tx)
356 .await?;
357
358 let update_raw = format!(
359 "UPDATE memory_tree
360 SET parent_id = ?, consolidated_at = {now}
361 WHERE id = ? AND parent_id IS NULL"
362 );
363 let update_sql = zeph_db::rewrite_placeholders(&update_raw);
364 for &child_id in child_ids {
365 zeph_db::query(sqlx::AssertSqlSafe(update_sql.as_str()))
366 .bind(parent_id)
367 .bind(child_id)
368 .execute(&mut *tx)
369 .await?;
370 }
371
372 tx.commit().await?;
373 Ok(parent_id)
374 }
375
376 pub async fn increment_tree_consolidation_count(&self) -> Result<(), MemoryError> {
382 let now = <ActiveDialect as zeph_db::dialect::Dialect>::NOW;
383 let raw = format!(
384 "UPDATE memory_tree_meta
385 SET total_consolidations = total_consolidations + 1,
386 last_consolidation_at = {now},
387 updated_at = {now}
388 WHERE id = 1"
389 );
390 let query_sql = zeph_db::rewrite_placeholders(&raw);
391 query(sqlx::AssertSqlSafe(query_sql))
392 .execute(self.pool())
393 .await?;
394
395 Ok(())
396 }
397
398 pub async fn count_tree_nodes(&self) -> Result<i64, MemoryError> {
404 let count: i64 = query_scalar(sql!("SELECT COUNT(*) FROM memory_tree"))
405 .fetch_one(self.pool())
406 .await?;
407
408 Ok(count)
409 }
410}
411
412#[cfg(test)]
413mod tests {
414 use super::*;
415
416 async fn make_store() -> SqliteStore {
417 SqliteStore::with_pool_size(":memory:", 1)
418 .await
419 .expect("in-memory store")
420 }
421
422 #[tokio::test]
423 async fn insert_leaf_and_count() {
424 let store = make_store().await;
425 let id = store
426 .insert_tree_leaf("remember this fact", 10)
427 .await
428 .expect("insert leaf");
429 assert!(id > 0);
430 assert_eq!(store.count_tree_nodes().await.expect("count"), 1);
431 }
432
433 #[tokio::test]
434 async fn load_unconsolidated_leaves_excludes_parented() {
435 let store = make_store().await;
436 let leaf1 = store.insert_tree_leaf("leaf one", 5).await.expect("leaf1");
437 let leaf2 = store.insert_tree_leaf("leaf two", 5).await.expect("leaf2");
438
439 let parent_id = store
441 .insert_tree_node(1, None, "summary of leaf1 and leaf2", "[]", 10)
442 .await
443 .expect("parent");
444 store
445 .mark_nodes_consolidated(&[leaf1, leaf2], parent_id)
446 .await
447 .expect("mark consolidated");
448
449 let leaves = store
451 .load_tree_leaves_unconsolidated(10)
452 .await
453 .expect("load");
454 assert!(
455 leaves.is_empty(),
456 "consolidated leaves must not appear in unconsolidated query"
457 );
458 }
459
460 #[tokio::test]
461 async fn mark_nodes_consolidated_is_per_cluster_transaction() {
462 let store = make_store().await;
463 let leaf1 = store.insert_tree_leaf("a", 1).await.expect("l1");
464 let leaf2 = store.insert_tree_leaf("b", 1).await.expect("l2");
465 let parent = store
466 .insert_tree_node(1, None, "ab summary", "[]", 2)
467 .await
468 .expect("parent");
469
470 store
471 .mark_nodes_consolidated(&[leaf1, leaf2], parent)
472 .await
473 .expect("mark");
474
475 let rows: Vec<MemoryTreeRow> = zeph_db::query_as(zeph_db::sql!(
477 "SELECT id, level, parent_id, content, source_ids, token_count,
478 consolidated_at, created_at
479 FROM memory_tree WHERE level = 0"
480 ))
481 .fetch_all(store.pool())
482 .await
483 .expect("fetch");
484
485 assert!(rows.iter().all(|r| r.parent_id == Some(parent)));
486 }
487
488 #[tokio::test]
489 async fn traverse_tree_up_returns_path() {
490 let store = make_store().await;
491 let leaf = store.insert_tree_leaf("leaf", 1).await.expect("leaf");
492 let mid = store
493 .insert_tree_node(1, None, "mid", "[]", 2)
494 .await
495 .expect("mid");
496 store
497 .mark_nodes_consolidated(&[leaf], mid)
498 .await
499 .expect("mark l→m");
500
501 let path = store.traverse_tree_up(leaf, 3).await.expect("traverse");
502 assert_eq!(path.len(), 2, "leaf + mid parent");
503 assert_eq!(path[0].id, leaf);
504 assert_eq!(path[1].id, mid);
505 }
506
507 #[tokio::test]
508 async fn mark_nodes_consolidated_empty_slice_is_noop() {
509 let store = make_store().await;
510 store.mark_nodes_consolidated(&[], 999).await.expect("noop");
512 }
513
514 #[tokio::test]
515 async fn load_tree_leaves_unconsolidated_empty_tree_is_empty() {
516 let store = make_store().await;
517 let leaves = store
518 .load_tree_leaves_unconsolidated(10)
519 .await
520 .expect("load");
521 assert!(leaves.is_empty());
522 }
523
524 #[tokio::test]
525 async fn bump_consolidation_attempts_empty_slice_is_noop() {
526 let store = make_store().await;
527 store.bump_consolidation_attempts(&[]).await.expect("noop");
528 }
529
530 async fn backdate_last_attempted(store: &SqliteStore, ids: &[i64], timestamp: &str) {
533 for &id in ids {
534 zeph_db::query(zeph_db::sql!(
535 "UPDATE memory_tree SET last_attempted_at = ? WHERE id = ?"
536 ))
537 .bind(timestamp)
538 .bind(id)
539 .execute(store.pool())
540 .await
541 .expect("backdate last_attempted_at");
542 }
543 }
544
545 async fn backdate_created_at(store: &SqliteStore, ids: &[i64], timestamp: &str) {
547 for &id in ids {
548 zeph_db::query(zeph_db::sql!(
549 "UPDATE memory_tree SET created_at = ? WHERE id = ?"
550 ))
551 .bind(timestamp)
552 .bind(id)
553 .execute(store.pool())
554 .await
555 .expect("backdate created_at");
556 }
557 }
558
559 async fn load_ids(store: &SqliteStore, limit: usize) -> std::collections::HashSet<i64> {
560 store
561 .load_tree_leaves_unconsolidated(limit)
562 .await
563 .expect("load")
564 .into_iter()
565 .map(|r| r.id)
566 .collect()
567 }
568
569 #[tokio::test]
581 async fn bump_consolidation_attempts_prevents_permanent_starvation() {
582 let store = make_store().await;
583
584 let stuck_a = store.insert_tree_leaf("stuck a", 1).await.expect("stuck_a");
585 let stuck_b = store.insert_tree_leaf("stuck b", 1).await.expect("stuck_b");
586
587 store
590 .bump_consolidation_attempts(&[stuck_a, stuck_b])
591 .await
592 .expect("bump 1");
593 backdate_last_attempted(&store, &[stuck_a, stuck_b], "2000-01-01 00:00:01").await;
594
595 let waiting = store
601 .insert_tree_leaf("waiting, created after sweep 1's touch", 1)
602 .await
603 .expect("waiting");
604 backdate_created_at(&store, &[waiting], "2000-01-01 00:00:02").await;
605
606 let batch1 = load_ids(&store, 1).await;
607 assert!(
608 batch1 == std::collections::HashSet::from([stuck_a])
609 || batch1 == std::collections::HashSet::from([stuck_b]),
610 "a just-touched stuck leaf must still outrank a leaf created immediately \
611 afterward, got {batch1:?}"
612 );
613
614 store
619 .bump_consolidation_attempts(&[stuck_a, stuck_b])
620 .await
621 .expect("bump 2");
622 backdate_last_attempted(&store, &[stuck_a, stuck_b], "2000-01-01 00:00:03").await;
623
624 let batch2 = load_ids(&store, 1).await;
625 assert_eq!(
626 batch2,
627 std::collections::HashSet::from([waiting]),
628 "a leaf that predates the actively-cycling leaves' most recent touch must win once \
629 they are touched again without it ever being touched itself"
630 );
631 }
632
633 #[tokio::test]
639 async fn load_tree_leaves_unconsolidated_prioritizes_stale_touch_over_sustained_new_arrivals() {
640 let store = make_store().await;
641
642 let stuck = store
643 .insert_tree_leaf("stuck singleton", 1)
644 .await
645 .expect("stuck");
646 zeph_db::query(zeph_db::sql!(
649 "UPDATE memory_tree
650 SET consolidation_attempts = 5, last_attempted_at = '2000-01-01 00:00:00'
651 WHERE id = ?"
652 ))
653 .bind(stuck)
654 .execute(store.pool())
655 .await
656 .expect("backdate stuck leaf");
657
658 for i in 0..5 {
661 store
662 .insert_tree_leaf(&format!("fresh leaf {i}"), 1)
663 .await
664 .expect("fresh");
665 }
666
667 let batch = store
668 .load_tree_leaves_unconsolidated(1)
669 .await
670 .expect("load");
671 assert_eq!(batch.len(), 1);
672 assert_eq!(
673 batch[0].id, stuck,
674 "a leaf neglected since 2000 must outrank leaves created moments ago"
675 );
676 }
677
678 #[tokio::test]
683 async fn load_tree_leaves_unconsolidated_returns_created_at_order_even_when_priority_differs() {
684 let store = make_store().await;
685
686 let older = store
687 .insert_tree_leaf("older, but recently retried", 1)
688 .await
689 .expect("older");
690 let newer = store
691 .insert_tree_leaf("newer, never attempted", 1)
692 .await
693 .expect("newer");
694
695 store
698 .bump_consolidation_attempts(&[older])
699 .await
700 .expect("bump older");
701
702 let batch = store
703 .load_tree_leaves_unconsolidated(2)
704 .await
705 .expect("load");
706 let ids: Vec<i64> = batch.iter().map(|r| r.id).collect();
707 assert_eq!(
708 ids,
709 [older, newer],
710 "returned rows must stay in created_at ASC order regardless of selection priority"
711 );
712 }
713}