1use super::*;
4
5pub(super) fn open_queue_db<P: AsRef<std::path::Path>>(path: P) -> Result<Connection, AppError> {
20 let conn = Connection::open(path)?;
21 conn.pragma_update(None, "journal_mode", "wal")?;
22 conn.pragma_update(None, "busy_timeout", crate::constants::BUSY_TIMEOUT_MILLIS)?;
28 conn.execute_batch(
29 "CREATE TABLE IF NOT EXISTS queue (
30 id INTEGER PRIMARY KEY AUTOINCREMENT,
31 item_key TEXT NOT NULL UNIQUE,
32 item_type TEXT NOT NULL DEFAULT 'memory',
33 status TEXT NOT NULL DEFAULT 'pending',
34 memory_id INTEGER,
35 entity_id INTEGER,
36 entities INTEGER DEFAULT 0,
37 rels INTEGER DEFAULT 0,
38 error TEXT,
39 cost_usd REAL DEFAULT 0.0,
40 attempt INTEGER DEFAULT 0,
41 elapsed_ms INTEGER,
42 created_at TEXT DEFAULT (datetime('now')),
43 done_at TEXT
44 );
45 CREATE INDEX IF NOT EXISTS idx_enrich_queue_status ON queue(status);",
46 )?;
47 let mut has_error_class = false;
51 let mut has_next_retry_at = false;
52 let mut has_operation = false;
57 let mut has_finish_reason = false;
63 let mut has_input_tokens = false;
64 let mut has_output_tokens = false;
65 let mut has_claimed_at = false;
70 {
71 let mut stmt = conn.prepare("PRAGMA table_info(queue)")?;
72 let names = stmt.query_map([], |r| r.get::<_, String>(1))?;
73 for name in names {
74 match name?.as_str() {
75 "error_class" => has_error_class = true,
76 "next_retry_at" => has_next_retry_at = true,
77 "operation" => has_operation = true,
78 "finish_reason" => has_finish_reason = true,
79 "input_tokens" => has_input_tokens = true,
80 "output_tokens" => has_output_tokens = true,
81 "claimed_at" => has_claimed_at = true,
82 _ => {}
83 }
84 }
85 }
86 if !has_error_class {
87 conn.execute_batch("ALTER TABLE queue ADD COLUMN error_class TEXT")?;
88 }
89 if !has_next_retry_at {
90 conn.execute_batch("ALTER TABLE queue ADD COLUMN next_retry_at TEXT")?;
91 }
92 if !has_operation {
93 conn.execute_batch("ALTER TABLE queue ADD COLUMN operation TEXT")?;
94 }
95 if !has_finish_reason {
96 conn.execute_batch("ALTER TABLE queue ADD COLUMN finish_reason TEXT")?;
97 }
98 if !has_input_tokens {
99 conn.execute_batch("ALTER TABLE queue ADD COLUMN input_tokens INTEGER")?;
100 }
101 if !has_output_tokens {
102 conn.execute_batch("ALTER TABLE queue ADD COLUMN output_tokens INTEGER")?;
103 }
104 if !has_claimed_at {
105 conn.execute_batch("ALTER TABLE queue ADD COLUMN claimed_at INTEGER")?;
106 }
107 conn.execute_batch(
108 "CREATE INDEX IF NOT EXISTS idx_enrich_queue_eligible ON queue(status, next_retry_at);
109 CREATE INDEX IF NOT EXISTS idx_enrich_queue_operation ON queue(operation, status);
110 CREATE INDEX IF NOT EXISTS idx_enrich_queue_memory ON queue(memory_id)",
111 )?;
112 Ok(conn)
113}
114
115pub(super) fn enqueue_candidate(
127 queue_conn: &Connection,
128 main_conn: &Connection,
129 namespace: &str,
130 key: &str,
131 item_type: &str,
132 operation: &str,
133) {
134 let memory_id: Option<i64> = if item_type == "memory" {
135 main_conn
136 .query_row(
137 "SELECT id FROM memories WHERE namespace=?1 AND name=?2 AND deleted_at IS NULL",
138 rusqlite::params![namespace, key],
139 |r| r.get(0),
140 )
141 .ok()
142 } else {
143 None
144 };
145 if let Err(e) = queue_conn.execute(
146 "INSERT OR IGNORE INTO queue (item_key, item_type, status, operation, memory_id) \
147 VALUES (?1, ?2, 'pending', ?3, ?4)",
148 rusqlite::params![key, item_type, operation, memory_id],
149 ) {
150 tracing::warn!(target: "enrich", error = %e, "queue insert failed");
151 }
152}
153
154pub fn reset_stale_processing_claims(
160 conn: &Connection,
161 max_age_secs: u64,
162) -> Result<usize, AppError> {
163 let reset = conn.execute(
164 "UPDATE queue SET status='pending', claimed_at=NULL \
165 WHERE status='processing' AND claimed_at IS NOT NULL \
166 AND CAST(strftime('%s','now') AS INTEGER) - claimed_at > ?1",
167 rusqlite::params![max_age_secs as i64],
168 )?;
169 Ok(reset)
170}
171
172pub fn heartbeat(conn: &Connection, queue_id: i64) -> Result<(), AppError> {
176 conn.execute(
177 "UPDATE queue SET claimed_at = CAST(strftime('%s','now') AS INTEGER) WHERE id = ?1",
178 rusqlite::params![queue_id],
179 )?;
180 Ok(())
181}
182
183pub(super) fn skipped_item_keys(
191 conn: &Connection,
192 operation: &str,
193) -> Result<std::collections::HashSet<String>, AppError> {
194 let mut stmt = conn.prepare(
195 "SELECT item_key FROM queue WHERE status='skipped' AND (operation = ?1 OR operation IS NULL)",
196 )?;
197 let keys = stmt
198 .query_map(rusqlite::params![operation], |r| r.get::<_, String>(0))?
199 .collect::<Result<std::collections::HashSet<String>, _>>()?;
200 Ok(keys)
201}
202
203pub(super) fn item_type_for(operation: &EnrichOperation) -> &'static str {
207 match operation {
208 EnrichOperation::EntityDescriptions => "entity",
209 EnrichOperation::EntityConnect | EnrichOperation::CrossDomainBridges => "entity_pair",
212 _ => "memory",
213 }
214}
215
216pub(super) fn item_type_for_key(key: &str, default: &'static str) -> &'static str {
226 if key.starts_with("pair:") {
227 "entity_pair"
228 } else if key.starts_with("entity:") {
229 "entity"
230 } else if key.starts_with("chunk:") {
231 "chunk"
232 } else {
233 default
234 }
235}
236
237pub fn cleanup_queue_entry(db_path: &std::path::Path, memory_id: i64, name: &str) {
246 let queue_path = crate::paths::sidecar_path(db_path, ".enrich-queue.sqlite");
247 if !queue_path.exists() {
248 return;
249 }
250 match open_queue_db(&queue_path) {
251 Ok(conn) => {
252 if let Err(e) = conn.execute(
253 "DELETE FROM queue WHERE memory_id = ?1 OR item_key = ?2",
254 rusqlite::params![memory_id, name],
255 ) {
256 tracing::warn!(target: "enrich", error = %e, memory_id, "enrich-queue cleanup failed");
257 }
258 }
259 Err(e) => {
260 tracing::warn!(target: "enrich", error = %e, "enrich-queue cleanup skipped (open failed)");
261 }
262 }
263}
264
265pub(super) fn prune_dead_orphans(
276 queue_conn: &Connection,
277 main_conn: &Connection,
278 operation: &str,
279 namespace: &str,
280) -> Result<i64, AppError> {
281 let dead: Vec<(i64, String)> = {
282 let mut stmt = queue_conn.prepare(
283 "SELECT id, item_key FROM queue \
284 WHERE status='dead' AND item_type='memory' \
285 AND (operation = ?1 OR operation IS NULL) ORDER BY id",
286 )?;
287 let rows = stmt
288 .query_map(rusqlite::params![operation], |r| Ok((r.get(0)?, r.get(1)?)))?
289 .collect::<Result<Vec<_>, _>>()?;
290 rows
291 };
292 let mut pruned = 0_i64;
293 for (id, name) in dead {
294 let exists = main_conn
295 .query_row(
296 "SELECT 1 FROM memories WHERE namespace=?1 AND name=?2 AND deleted_at IS NULL",
297 rusqlite::params![namespace, name],
298 |_| Ok(()),
299 )
300 .is_ok();
301 if !exists {
302 queue_conn.execute("DELETE FROM queue WHERE id=?1", rusqlite::params![id])?;
303 pruned += 1;
304 }
305 }
306 if pruned > 0 {
307 let _ = queue_conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);");
308 }
309 Ok(pruned)
310}
311
312pub(super) fn prune_dead_entity_orphans(
318 queue_conn: &Connection,
319 operation: &str,
320) -> Result<i64, AppError> {
321 let pruned = queue_conn.execute(
322 "DELETE FROM queue \
323 WHERE status='dead' AND item_type='entity' \
324 AND (operation = ?1 OR operation IS NULL)",
325 rusqlite::params![operation],
326 )? as i64;
327 if pruned > 0 {
328 let _ = queue_conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);");
329 }
330 Ok(pruned)
331}
332
333#[derive(Debug, Serialize, schemars::JsonSchema)]
343pub struct EnrichStatus {
344 pub(super) status_report: bool,
345 pub(super) operation: String,
346 pub(super) namespace: String,
347 pub(super) unbound_backlog: usize,
348 pub(super) scan_backlog: i64,
355 pub(super) queue_pending: i64,
356 pub(super) queue_processing: i64,
357 pub(super) queue_done: i64,
358 pub(super) queue_failed: i64,
359 pub(super) queue_skipped: i64,
360 pub(super) queue_dead: i64,
361 pub(super) eligible_now: i64,
362 pub(super) waiting: i64,
363 pub(super) state: &'static str,
369 pub(super) waiting_items: Vec<WaitingItem>,
372}
373
374#[derive(Debug, Serialize, schemars::JsonSchema)]
376pub struct WaitingItem {
377 pub(super) item_key: String,
378 pub(super) attempt: i64,
379 pub(super) next_retry_at: Option<String>,
380 pub(super) error_class: Option<String>,
381}
382
383#[derive(Debug, Serialize, schemars::JsonSchema)]
385pub struct DeadItem {
386 pub(super) dead_item: bool,
387 pub(super) item_key: String,
388 pub(super) item_type: String,
389 pub(super) attempt: i64,
390 pub(super) error_class: Option<String>,
391 pub(super) error: Option<String>,
392 pub(super) finish_reason: Option<String>,
397 pub(super) input_tokens: Option<i64>,
399 pub(super) output_tokens: Option<i64>,
401}
402
403#[derive(Debug, Serialize, schemars::JsonSchema)]
405pub struct DeadSummary {
406 pub(super) summary: bool,
407 pub(super) operation: String,
408 pub(super) namespace: String,
409 pub(super) action: &'static str,
411 pub(super) dead_total: i64,
412 pub(super) requeued: i64,
413 pub(super) pruned: i64,
417}
418
419pub(super) fn classify_enrich_outcome(e: &AppError) -> crate::retry::AttemptOutcome {
430 use crate::retry::AttemptOutcome;
431 match e {
432 AppError::RateLimited { .. } | AppError::Timeout { .. } | AppError::DbBusy(_) => {
433 AttemptOutcome::Transient
434 }
435 AppError::EntityNotYetMaterialized { .. } => AttemptOutcome::Transient,
443 AppError::ProviderError { .. }
449 | AppError::NotFound(_)
450 | AppError::MemoryNotFound { .. }
451 | AppError::MemoryNotFoundById { .. } => AttemptOutcome::HardFailure,
452 AppError::Database(_) => {
456 if crate::storage::utils::is_sqlite_busy(e) {
457 AttemptOutcome::Transient
458 } else {
459 AttemptOutcome::HardFailure
460 }
461 }
462 AppError::Embedding(_) => AttemptOutcome::Transient,
473 _ => AttemptOutcome::HardFailure,
484 }
485}
486
487pub(super) fn record_item_failure(
500 queue_conn: &rusqlite::Connection,
501 queue_id: i64,
502 attempt: i64,
503 max_attempts: u32,
504 err: &AppError,
505) -> crate::retry::AttemptOutcome {
506 let outcome = classify_enrich_outcome(err);
507 let err_str = format!("{err}");
508 record_item_failure_typed(
509 queue_conn,
510 queue_id,
511 attempt,
512 max_attempts,
513 outcome,
514 &err_str,
515 None,
516 None,
517 None,
518 )
519}
520
521#[allow(clippy::too_many_arguments)]
530pub(super) fn record_item_failure_typed(
531 queue_conn: &rusqlite::Connection,
532 queue_id: i64,
533 attempt: i64,
534 max_attempts: u32,
535 outcome: crate::retry::AttemptOutcome,
536 err_str: &str,
537 finish_reason: Option<&str>,
538 input_tokens: Option<i64>,
539 output_tokens: Option<i64>,
540) -> crate::retry::AttemptOutcome {
541 use crate::retry::AttemptOutcome;
542 let error_class = match outcome {
543 AttemptOutcome::Transient => "transient",
544 AttemptOutcome::HardFailure => "permanent",
545 AttemptOutcome::Success => "success",
546 };
547
548 let terminal = matches!(outcome, AttemptOutcome::HardFailure) || attempt >= max_attempts as i64;
549 if terminal {
550 let _ = queue_conn.execute(
551 "UPDATE queue SET status='dead', error=?1, error_class=?2, done_at=datetime('now'), \
552 finish_reason=?3, input_tokens=?4, output_tokens=?5 WHERE id=?6",
553 rusqlite::params![
554 err_str,
555 error_class,
556 finish_reason,
557 input_tokens,
558 output_tokens,
559 queue_id
560 ],
561 );
562 } else {
563 let delay = crate::retry::compute_delay(
564 &crate::retry::RetryConfig::llm_rate_limit(),
565 attempt.max(0) as u32,
566 );
567 let secs = delay.as_secs().max(1);
568 let modifier = format!("+{secs} seconds");
569 let _ = queue_conn.execute(
570 "UPDATE queue SET status='pending', error=?1, error_class=?2, next_retry_at=datetime('now', ?3), \
571 finish_reason=?4, input_tokens=?5, output_tokens=?6 WHERE id=?7",
572 rusqlite::params![
573 err_str,
574 error_class,
575 modifier,
576 finish_reason,
577 input_tokens,
578 output_tokens,
579 queue_id
580 ],
581 );
582 }
583 outcome
584}
585
586pub(super) enum DequeueOutcome {
594 Claimed((i64, String, String, i64)),
595 Empty,
596}
597
598pub(super) fn dequeue_next_pending(
599 queue_conn: &rusqlite::Connection,
600 backoff_clause: &str,
601) -> Result<DequeueOutcome, AppError> {
602 let dequeue_sql = format!(
603 "UPDATE queue SET status='processing', attempt=attempt+1, \
604 claimed_at=CAST(strftime('%s','now') AS INTEGER) \
605 WHERE id = (SELECT id FROM queue WHERE status='pending' {backoff_clause} \
606 ORDER BY id LIMIT 1) \
607 RETURNING id, item_key, item_type, attempt"
608 );
609 match queue_conn.query_row(&dequeue_sql, [], |row| {
610 Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?))
611 }) {
612 Ok(claimed) => Ok(DequeueOutcome::Claimed(claimed)),
613 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(DequeueOutcome::Empty),
614 Err(e) => Err(AppError::Database(e)),
615 }
616}
617
618#[cfg(test)]
623mod tests {
624 use super::*;
625
626 fn open_test_db() -> Connection {
627 let conn = Connection::open_in_memory().expect("in-memory db");
628 conn.execute_batch(
629 "CREATE TABLE memories (
630 id INTEGER PRIMARY KEY AUTOINCREMENT,
631 namespace TEXT NOT NULL DEFAULT 'global',
632 name TEXT NOT NULL,
633 type TEXT NOT NULL DEFAULT 'note',
634 description TEXT NOT NULL DEFAULT '',
635 body TEXT NOT NULL DEFAULT '',
636 body_hash TEXT NOT NULL DEFAULT '',
637 session_id TEXT,
638 source TEXT NOT NULL DEFAULT 'agent',
639 metadata TEXT NOT NULL DEFAULT '{}',
640 created_at INTEGER NOT NULL DEFAULT (unixepoch()),
641 updated_at INTEGER NOT NULL DEFAULT (unixepoch()),
642 deleted_at INTEGER,
643 UNIQUE(namespace, name)
644 );",
645 )
646 .expect("schema creation must succeed");
647 conn
648 }
649
650 fn open_temp_queue() -> (Connection, String) {
651 let path = format!(
652 "/tmp/test-enrich-dl-{}-{}.sqlite",
653 std::process::id(),
654 fastrand::u64(..)
655 );
656 let conn = open_queue_db(&path).expect("queue db must open");
657 (conn, path)
658 }
659
660 fn insert_pending(conn: &Connection, key: &str) -> i64 {
661 conn.execute(
662 "INSERT INTO queue (item_key, item_type, status) VALUES (?1, 'memory', 'pending')",
663 rusqlite::params![key],
664 )
665 .unwrap();
666 conn.last_insert_rowid()
667 }
668
669 #[test]
670 fn queue_db_schema_creates_correctly() {
671 let tmp_path = format!("/tmp/test-enrich-queue-{}.sqlite", std::process::id());
672 let conn = open_queue_db(&tmp_path).expect("queue db must open");
673 let count: i64 = conn
674 .query_row("SELECT COUNT(*) FROM queue", [], |r| r.get(0))
675 .unwrap();
676 assert_eq!(count, 0);
677 let _ = std::fs::remove_file(&tmp_path);
678 }
679
680 #[test]
681 fn classify_rate_limit_is_transient() {
682 let e = AppError::RateLimited {
683 detail: "429".into(),
684 };
685 assert_eq!(
686 classify_enrich_outcome(&e),
687 crate::retry::AttemptOutcome::Transient
688 );
689 }
690
691 #[test]
692 fn classify_timeout_and_dbbusy_are_transient() {
693 let t = AppError::Timeout {
694 operation: "judge".into(),
695 duration_secs: 30,
696 };
697 let b = AppError::DbBusy("locked".into());
698 assert_eq!(
699 classify_enrich_outcome(&t),
700 crate::retry::AttemptOutcome::Transient
701 );
702 assert_eq!(
703 classify_enrich_outcome(&b),
704 crate::retry::AttemptOutcome::Transient
705 );
706 }
707
708 #[test]
709 fn classify_validation_and_parse_are_hard_failure() {
710 let v = AppError::Validation("failed to parse entities array: bad".into());
711 assert_eq!(
712 classify_enrich_outcome(&v),
713 crate::retry::AttemptOutcome::HardFailure
714 );
715 }
716
717 #[test]
718 fn open_queue_db_alter_is_idempotent() {
719 let path = format!(
720 "/tmp/test-enrich-idem-{}-{}.sqlite",
721 std::process::id(),
722 fastrand::u64(..)
723 );
724 let _ = open_queue_db(&path).expect("first open");
725 let conn = open_queue_db(&path).expect("second open is idempotent");
726 let cols: Vec<String> = {
727 let mut stmt = conn.prepare("PRAGMA table_info(queue)").unwrap();
728 stmt.query_map([], |r| r.get::<_, String>(1))
729 .unwrap()
730 .collect::<Result<Vec<_>, _>>()
731 .unwrap()
732 };
733 assert!(cols.iter().any(|c| c == "error_class"));
734 assert!(cols.iter().any(|c| c == "next_retry_at"));
735 let _ = std::fs::remove_file(&path);
736 }
737
738 #[test]
739 fn record_item_failure_hard_marks_dead() {
740 let (conn, path) = open_temp_queue();
741 let id = insert_pending(&conn, "mem-hard");
742 let outcome = record_item_failure(
743 &conn,
744 id,
745 1,
746 5,
747 &AppError::Validation("invalid body".into()),
748 );
749 assert_eq!(outcome, crate::retry::AttemptOutcome::HardFailure);
750 let status: String = conn
751 .query_row(
752 "SELECT status FROM queue WHERE id=?1",
753 rusqlite::params![id],
754 |r| r.get(0),
755 )
756 .unwrap();
757 assert_eq!(status, "dead");
758 let _ = std::fs::remove_file(&path);
759 }
760
761 #[test]
762 fn record_item_failure_transient_reschedules_pending() {
763 let (conn, path) = open_temp_queue();
764 let id = insert_pending(&conn, "mem-transient");
765 let outcome = record_item_failure(
766 &conn,
767 id,
768 1,
769 5,
770 &AppError::RateLimited {
771 detail: "429".into(),
772 },
773 );
774 assert_eq!(outcome, crate::retry::AttemptOutcome::Transient);
775 let (status, future): (String, i64) = conn
776 .query_row(
777 "SELECT status, (next_retry_at > datetime('now')) FROM queue WHERE id=?1",
778 rusqlite::params![id],
779 |r| Ok((r.get(0)?, r.get(1)?)),
780 )
781 .unwrap();
782 assert_eq!(status, "pending");
783 assert_eq!(future, 1, "next_retry_at must be in the future");
784 let _ = std::fs::remove_file(&path);
785 }
786
787 #[test]
788 fn record_item_failure_transient_at_cap_marks_dead() {
789 let (conn, path) = open_temp_queue();
790 let id = insert_pending(&conn, "mem-cap");
791 let outcome = record_item_failure(
792 &conn,
793 id,
794 5,
795 5,
796 &AppError::RateLimited {
797 detail: "429".into(),
798 },
799 );
800 assert_eq!(outcome, crate::retry::AttemptOutcome::Transient);
801 let status: String = conn
802 .query_row(
803 "SELECT status FROM queue WHERE id=?1",
804 rusqlite::params![id],
805 |r| r.get(0),
806 )
807 .unwrap();
808 assert_eq!(status, "dead");
809 let _ = std::fs::remove_file(&path);
810 }
811
812 #[test]
813 fn dequeue_skips_future_retry_and_dead() {
814 let (conn, path) = open_temp_queue();
815 let eligible = insert_pending(&conn, "mem-eligible");
816 let waiting = insert_pending(&conn, "mem-waiting");
817 conn.execute(
818 "UPDATE queue SET next_retry_at=datetime('now', '+3600 seconds') WHERE id=?1",
819 rusqlite::params![waiting],
820 )
821 .unwrap();
822 let dead = insert_pending(&conn, "mem-dead");
823 conn.execute(
824 "UPDATE queue SET status='dead' WHERE id=?1",
825 rusqlite::params![dead],
826 )
827 .unwrap();
828
829 let claimed: Option<i64> = conn
830 .query_row(
831 "UPDATE queue SET status='processing', attempt=attempt+1 \
832 WHERE id = (SELECT id FROM queue WHERE status='pending' \
833 AND (next_retry_at IS NULL OR next_retry_at <= datetime('now')) \
834 ORDER BY id LIMIT 1) \
835 RETURNING id",
836 [],
837 |r| r.get(0),
838 )
839 .ok();
840 assert_eq!(claimed, Some(eligible));
841
842 let second: Option<i64> = conn
843 .query_row(
844 "UPDATE queue SET status='processing', attempt=attempt+1 \
845 WHERE id = (SELECT id FROM queue WHERE status='pending' \
846 AND (next_retry_at IS NULL OR next_retry_at <= datetime('now')) \
847 ORDER BY id LIMIT 1) \
848 RETURNING id",
849 [],
850 |r| r.get(0),
851 )
852 .ok();
853 assert_eq!(second, None);
854 let _ = std::fs::remove_file(&path);
855 }
856
857 #[test]
858 fn classify_validation_never_infers_transience_from_message() {
859 for msg in [
867 "model 'x' returned non-object JSON after repair (got string)",
868 "model 'x' returned content that could not be parsed even after JSON repair",
869 "model 'x' returned no structured content",
870 "LLM result missing 'description' field",
871 "LLM result missing 'enriched_body' field",
872 ] {
873 assert_eq!(
874 classify_enrich_outcome(&AppError::Validation(msg.into())),
875 crate::retry::AttemptOutcome::HardFailure,
876 "expected hard failure for: {msg}"
877 );
878 }
879 }
880
881 #[test]
882 fn classify_embedding_error_is_transient_floor() {
883 assert_eq!(
884 classify_enrich_outcome(&AppError::Embedding("dimension mismatch".into())),
885 crate::retry::AttemptOutcome::Transient
886 );
887 }
888
889 #[test]
892 fn classify_entity_not_yet_materialized_is_transient() {
893 assert_eq!(
894 classify_enrich_outcome(&AppError::EntityNotYetMaterialized {
895 name: "acme".into(),
896 namespace: "global".into(),
897 }),
898 crate::retry::AttemptOutcome::Transient
899 );
900 }
901
902 #[test]
903 fn classify_memory_absence_stays_hard_failure() {
904 assert_eq!(
905 classify_enrich_outcome(&AppError::MemoryNotFound {
906 name: "mem-x".into(),
907 namespace: "global".into(),
908 }),
909 crate::retry::AttemptOutcome::HardFailure
910 );
911 assert_eq!(
912 classify_enrich_outcome(&AppError::MemoryNotFoundById { id: 42 }),
913 crate::retry::AttemptOutcome::HardFailure
914 );
915 assert_eq!(
916 classify_enrich_outcome(&AppError::NotFound("gone".into())),
917 crate::retry::AttemptOutcome::HardFailure
918 );
919 }
920
921 #[test]
922 fn classify_database_busy_is_transient_non_busy_is_hard() {
923 let busy = AppError::Database(rusqlite::Error::SqliteFailure(
924 rusqlite::ffi::Error::new(rusqlite::ffi::SQLITE_BUSY),
925 Some("database is locked".into()),
926 ));
927 assert_eq!(
928 classify_enrich_outcome(&busy),
929 crate::retry::AttemptOutcome::Transient
930 );
931 let constraint = AppError::Database(rusqlite::Error::SqliteFailure(
932 rusqlite::ffi::Error::new(rusqlite::ffi::SQLITE_CONSTRAINT),
933 Some("UNIQUE constraint failed".into()),
934 ));
935 assert_eq!(
936 classify_enrich_outcome(&constraint),
937 crate::retry::AttemptOutcome::HardFailure
938 );
939 }
940
941 #[test]
942 fn record_item_failure_typed_persists_diagnostics_on_dead_letter() {
943 let (conn, path) = open_temp_queue();
944 let id = insert_pending(&conn, "mem-diag");
945 let outcome = record_item_failure_typed(
946 &conn,
947 id,
948 1,
949 5,
950 crate::retry::AttemptOutcome::HardFailure,
951 "truncated response",
952 Some("length"),
953 Some(120),
954 Some(4096),
955 );
956 assert_eq!(outcome, crate::retry::AttemptOutcome::HardFailure);
957 let (status, finish_reason, input_tokens, output_tokens): (
958 String,
959 Option<String>,
960 Option<i64>,
961 Option<i64>,
962 ) = conn
963 .query_row(
964 "SELECT status, finish_reason, input_tokens, output_tokens FROM queue WHERE id=?1",
965 rusqlite::params![id],
966 |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?)),
967 )
968 .unwrap();
969 assert_eq!(status, "dead");
970 assert_eq!(finish_reason.as_deref(), Some("length"));
971 assert_eq!(input_tokens, Some(120));
972 assert_eq!(output_tokens, Some(4096));
973 let _ = std::fs::remove_file(&path);
974 }
975
976 #[test]
977 fn record_item_failure_typed_reschedules_transient_below_max_attempts() {
978 let (conn, path) = open_temp_queue();
984 let id = insert_pending(&conn, "mem-retry");
985 let outcome = record_item_failure_typed(
986 &conn,
987 id,
988 1,
989 5,
990 crate::retry::AttemptOutcome::Transient,
991 "truncated response",
992 Some("length"),
993 Some(120),
994 Some(64),
995 );
996 assert_eq!(outcome, crate::retry::AttemptOutcome::Transient);
997 let (status, error_class, finish_reason, next_retry_at): (
998 String,
999 String,
1000 Option<String>,
1001 Option<String>,
1002 ) = conn
1003 .query_row(
1004 "SELECT status, error_class, finish_reason, next_retry_at FROM queue WHERE id=?1",
1005 rusqlite::params![id],
1006 |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?)),
1007 )
1008 .unwrap();
1009 assert_eq!(status, "pending");
1010 assert_eq!(error_class, "transient");
1011 assert_eq!(finish_reason.as_deref(), Some("length"));
1012 assert!(
1013 next_retry_at.is_some(),
1014 "a rescheduled item must carry a next_retry_at"
1015 );
1016 let _ = std::fs::remove_file(&path);
1017 }
1018
1019 #[test]
1029 fn with_busy_retry_bounds_dequeue_under_sustained_contention() {
1030 let (conn, path) = open_temp_queue();
1031 insert_pending(&conn, "mem-busy");
1032 conn.pragma_update(None, "busy_timeout", 0i64)
1033 .expect("busy_timeout override must succeed");
1034
1035 let blocker = Connection::open(&path).expect("blocker connection must open");
1039 blocker
1040 .execute_batch("BEGIN EXCLUSIVE;")
1041 .expect("exclusive lock must be acquired");
1042
1043 let calls = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
1044 let calls_clone = std::sync::Arc::clone(&calls);
1045 let result: Result<DequeueOutcome, AppError> =
1046 crate::storage::utils::with_busy_retry(|| {
1047 calls_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1048 dequeue_next_pending(&conn, "")
1049 });
1050
1051 assert!(
1052 matches!(result, Err(AppError::DbBusy(_))),
1053 "sustained SQLITE_BUSY must convert to DbBusy, not hang or silently report Empty"
1054 );
1055 assert_eq!(
1056 calls.load(std::sync::atomic::Ordering::SeqCst),
1057 crate::constants::MAX_SQLITE_BUSY_RETRIES,
1058 "must attempt exactly MAX_SQLITE_BUSY_RETRIES times, never retry unbounded"
1059 );
1060
1061 blocker
1062 .execute_batch("ROLLBACK;")
1063 .expect("releasing the exclusive lock must succeed");
1064 let _ = std::fs::remove_file(&path);
1065 }
1066
1067 #[test]
1068 fn dequeue_next_pending_distinguishes_empty_from_claimed() {
1069 let (conn, path) = open_temp_queue();
1070 let id = insert_pending(&conn, "mem-dequeue");
1071 let claimed = dequeue_next_pending(&conn, "").expect("dequeue must succeed");
1072 match claimed {
1073 DequeueOutcome::Claimed((claimed_id, key, _, _)) => {
1074 assert_eq!(claimed_id, id);
1075 assert_eq!(key, "mem-dequeue");
1076 }
1077 DequeueOutcome::Empty => panic!("expected a claimed row"),
1078 }
1079 let empty = dequeue_next_pending(&conn, "").expect("dequeue must succeed");
1080 assert!(matches!(empty, DequeueOutcome::Empty));
1081 let _ = std::fs::remove_file(&path);
1082 }
1083
1084 #[test]
1085 fn classify_provider_error_and_not_found_are_hard() {
1086 assert_eq!(
1087 classify_enrich_outcome(&AppError::ProviderError {
1088 code: "400".into(),
1089 message: "context length exceeded".into(),
1090 }),
1091 crate::retry::AttemptOutcome::HardFailure
1092 );
1093 assert_eq!(
1094 classify_enrich_outcome(&AppError::NotFound("memory 'gone' not found".into())),
1095 crate::retry::AttemptOutcome::HardFailure
1096 );
1097 }
1098
1099 #[test]
1100 fn open_queue_db_migrates_operation_column() {
1101 let (conn, path) = open_temp_queue();
1102 drop(conn);
1103 let conn = open_queue_db(&path).expect("second open is idempotent");
1104 let cols: Vec<String> = {
1105 let mut stmt = conn.prepare("PRAGMA table_info(queue)").unwrap();
1106 stmt.query_map([], |r| r.get::<_, String>(1))
1107 .unwrap()
1108 .collect::<Result<Vec<_>, _>>()
1109 .unwrap()
1110 };
1111 assert!(cols.iter().any(|c| c == "operation"));
1112 assert!(cols.iter().any(|c| c == "memory_id"));
1113 let _ = std::fs::remove_file(&path);
1114 }
1115
1116 #[test]
1117 fn enqueue_candidate_tags_operation_and_memory_id() {
1118 let main = open_test_db();
1119 main.execute(
1120 "INSERT INTO memories (namespace, name, body) VALUES ('global', 'mem-x', 'body')",
1121 [],
1122 )
1123 .unwrap();
1124 let mem_id: i64 = main
1125 .query_row("SELECT id FROM memories WHERE name='mem-x'", [], |r| {
1126 r.get(0)
1127 })
1128 .unwrap();
1129 let (queue, path) = open_temp_queue();
1130 enqueue_candidate(&queue, &main, "global", "mem-x", "memory", "MemoryBindings");
1131 let (op, mid): (String, i64) = queue
1132 .query_row(
1133 "SELECT operation, memory_id FROM queue WHERE item_key='mem-x'",
1134 [],
1135 |r| Ok((r.get(0)?, r.get(1)?)),
1136 )
1137 .unwrap();
1138 assert_eq!(op, "MemoryBindings");
1139 assert_eq!(mid, mem_id);
1140 let _ = std::fs::remove_file(&path);
1141 }
1142
1143 #[test]
1144 fn requeue_dead_resurrects_dead_rows() {
1145 let (conn, path) = open_temp_queue();
1146 conn.execute(
1147 "INSERT INTO queue (item_key, item_type, status, operation, attempt, error, error_class, next_retry_at) \
1148 VALUES ('mem-dead', 'memory', 'dead', 'MemoryBindings', 8, 'boom', 'permanent', datetime('now'))",
1149 [],
1150 )
1151 .unwrap();
1152 let n = conn
1153 .execute(
1154 "UPDATE queue SET status='pending', attempt=0, next_retry_at=NULL, \
1155 error=NULL, error_class=NULL \
1156 WHERE status='dead' AND (operation = ?1 OR operation IS NULL)",
1157 rusqlite::params!["MemoryBindings"],
1158 )
1159 .unwrap();
1160 assert_eq!(n, 1);
1161 let (status, attempt, nra): (String, i64, Option<String>) = conn
1162 .query_row(
1163 "SELECT status, attempt, next_retry_at FROM queue WHERE item_key='mem-dead'",
1164 [],
1165 |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)),
1166 )
1167 .unwrap();
1168 assert_eq!(status, "pending");
1169 assert_eq!(attempt, 0);
1170 assert!(nra.is_none());
1171 let _ = std::fs::remove_file(&path);
1172 }
1173
1174 #[test]
1175 fn skipped_item_keys_excludes_only_skipped_for_operation() {
1176 let (conn, path) = open_temp_queue();
1181 conn.execute(
1182 "INSERT INTO queue (item_key, item_type, status, operation) VALUES ('mem-vetoed', 'memory', 'skipped', 'BodyEnrich')",
1183 [],
1184 )
1185 .unwrap();
1186 conn.execute(
1187 "INSERT INTO queue (item_key, item_type, status, operation) VALUES ('mem-pending', 'memory', 'pending', 'BodyEnrich')",
1188 [],
1189 )
1190 .unwrap();
1191 conn.execute(
1192 "INSERT INTO queue (item_key, item_type, status, operation) VALUES ('mem-other-op', 'memory', 'skipped', 'MemoryBindings')",
1193 [],
1194 )
1195 .unwrap();
1196 let keys = skipped_item_keys(&conn, "BodyEnrich").unwrap();
1197 assert!(
1198 keys.contains("mem-vetoed"),
1199 "vetoed BodyEnrich item must be excluded from scan"
1200 );
1201 assert!(
1202 !keys.contains("mem-pending"),
1203 "pending item is still actionable"
1204 );
1205 assert!(
1206 !keys.contains("mem-other-op"),
1207 "skipped item from another operation must not leak"
1208 );
1209 assert_eq!(keys.len(), 1);
1210 let _ = std::fs::remove_file(&path);
1211 }
1212
1213 #[test]
1214 fn cascade_cleanup_delete_targets_memory_id_and_name() {
1215 let (conn, path) = open_temp_queue();
1216 conn.execute(
1217 "INSERT INTO queue (item_key, item_type, status, memory_id) VALUES ('by-id', 'memory', 'done', 42)",
1218 [],
1219 )
1220 .unwrap();
1221 conn.execute(
1222 "INSERT INTO queue (item_key, item_type, status) VALUES ('by-name', 'memory', 'pending')",
1223 [],
1224 )
1225 .unwrap();
1226 let removed = conn
1227 .execute(
1228 "DELETE FROM queue WHERE memory_id = ?1 OR item_key = ?2",
1229 rusqlite::params![42_i64, "by-name"],
1230 )
1231 .unwrap();
1232 assert_eq!(removed, 2);
1233 let remaining: i64 = conn
1234 .query_row("SELECT COUNT(*) FROM queue", [], |r| r.get(0))
1235 .unwrap();
1236 assert_eq!(remaining, 0);
1237 let _ = std::fs::remove_file(&path);
1238 }
1239
1240 #[test]
1241 fn item_type_for_maps_entity_and_memory() {
1242 assert_eq!(
1243 item_type_for(&EnrichOperation::EntityDescriptions),
1244 "entity"
1245 );
1246 assert_eq!(item_type_for(&EnrichOperation::MemoryBindings), "memory");
1247 assert_eq!(item_type_for(&EnrichOperation::AugmentBindings), "memory");
1248 assert_eq!(item_type_for(&EnrichOperation::BodyExtract), "memory");
1249 assert_eq!(
1250 item_type_for(&EnrichOperation::EntityConnect),
1251 "entity_pair"
1252 );
1253 assert_eq!(
1254 item_type_for(&EnrichOperation::CrossDomainBridges),
1255 "entity_pair"
1256 );
1257 }
1258
1259 #[test]
1262 fn item_type_for_key_honours_reembed_prefixes() {
1263 assert_eq!(item_type_for_key("plain-memory-name", "memory"), "memory");
1264 assert_eq!(
1265 item_type_for_key("entity:tokio-runtime", "memory"),
1266 "entity"
1267 );
1268 assert_eq!(item_type_for_key("chunk:42", "memory"), "chunk");
1269 assert_eq!(item_type_for_key("some-entity", "entity"), "entity");
1270 assert_eq!(item_type_for_key("pair:1:2", "memory"), "entity_pair");
1272 }
1273
1274 #[test]
1275 fn prune_dead_orphans_removes_only_orphan_memory_rows() {
1276 let main = open_test_db();
1277 main.execute(
1279 "INSERT INTO memories (namespace, name, body) VALUES ('global', 'alive', 'b')",
1280 [],
1281 )
1282 .unwrap();
1283 let (queue, path) = open_temp_queue();
1284 queue
1286 .execute(
1287 "INSERT INTO queue (item_key, item_type, status, operation, error_class) \
1288 VALUES ('gone', 'memory', 'dead', 'MemoryBindings', 'permanent')",
1289 [],
1290 )
1291 .unwrap();
1292 queue
1294 .execute(
1295 "INSERT INTO queue (item_key, item_type, status, operation, error_class) \
1296 VALUES ('alive', 'memory', 'dead', 'MemoryBindings', 'permanent')",
1297 [],
1298 )
1299 .unwrap();
1300 queue
1302 .execute(
1303 "INSERT INTO queue (item_key, item_type, status, operation) \
1304 VALUES ('some-entity', 'entity', 'dead', 'EntityDescriptions')",
1305 [],
1306 )
1307 .unwrap();
1308
1309 let pruned = prune_dead_orphans(&queue, &main, "MemoryBindings", "global").unwrap();
1310 assert_eq!(pruned, 1, "only the orphan memory row is pruned");
1311
1312 let remaining: Vec<String> = {
1313 let mut stmt = queue
1314 .prepare("SELECT item_key FROM queue ORDER BY item_key")
1315 .unwrap();
1316 stmt.query_map([], |r| r.get::<_, String>(0))
1317 .unwrap()
1318 .collect::<Result<Vec<_>, _>>()
1319 .unwrap()
1320 };
1321 assert_eq!(remaining, vec!["alive", "some-entity"]);
1322 let _ = std::fs::remove_file(&path);
1323 }
1324
1325 #[test]
1326 fn prune_dead_entity_orphans_removes_only_entity_dead_rows() {
1327 let (queue, path) = open_temp_queue();
1328 queue
1330 .execute(
1331 "INSERT INTO queue (item_key, item_type, status, operation, error_class) \
1332 VALUES ('entity:foo', 'entity', 'dead', 'ReEmbed', 'permanent')",
1333 [],
1334 )
1335 .unwrap();
1336 queue
1338 .execute(
1339 "INSERT INTO queue (item_key, item_type, status, operation, error_class) \
1340 VALUES ('mem-dead', 'memory', 'dead', 'MemoryBindings', 'permanent')",
1341 [],
1342 )
1343 .unwrap();
1344 queue
1346 .execute(
1347 "INSERT INTO queue (item_key, item_type, status, operation) \
1348 VALUES ('entity:bar', 'entity', 'pending', 'ReEmbed')",
1349 [],
1350 )
1351 .unwrap();
1352
1353 let pruned = prune_dead_entity_orphans(&queue, "ReEmbed").unwrap();
1354 assert_eq!(pruned, 1, "only the entity dead row is pruned");
1355
1356 let remaining: Vec<String> = {
1357 let mut stmt = queue
1358 .prepare("SELECT item_key FROM queue ORDER BY item_key")
1359 .unwrap();
1360 stmt.query_map([], |r| r.get::<_, String>(0))
1361 .unwrap()
1362 .collect::<Result<Vec<_>, _>>()
1363 .unwrap()
1364 };
1365 assert_eq!(remaining, vec!["entity:bar", "mem-dead"]);
1366 let _ = std::fs::remove_file(&path);
1367 }
1368
1369 #[test]
1372 fn stale_processing_claim_is_reset_after_threshold() {
1373 let (conn, path) = open_temp_queue();
1374 let id = insert_pending(&conn, "mem-stale");
1375 conn.execute(
1377 "UPDATE queue SET status='processing', claimed_at = CAST(strftime('%s','now') AS INTEGER) - 7200 WHERE id=?1",
1378 rusqlite::params![id],
1379 )
1380 .unwrap();
1381 let reset = reset_stale_processing_claims(&conn, 1800).unwrap();
1382 assert_eq!(reset, 1, "a stale claim older than the threshold is reset");
1383 let (status, claimed_at): (String, Option<i64>) = conn
1384 .query_row(
1385 "SELECT status, claimed_at FROM queue WHERE id=?1",
1386 rusqlite::params![id],
1387 |r| Ok((r.get(0)?, r.get(1)?)),
1388 )
1389 .unwrap();
1390 assert_eq!(status, "pending");
1391 assert!(claimed_at.is_none(), "claimed_at must be cleared on reset");
1392 let _ = std::fs::remove_file(&path);
1393 }
1394
1395 #[test]
1398 fn fresh_processing_claim_is_preserved() {
1399 let (conn, path) = open_temp_queue();
1400 let id = insert_pending(&conn, "mem-fresh");
1401 conn.execute(
1403 "UPDATE queue SET status='processing', claimed_at = CAST(strftime('%s','now') AS INTEGER) WHERE id=?1",
1404 rusqlite::params![id],
1405 )
1406 .unwrap();
1407 let reset = reset_stale_processing_claims(&conn, 1800).unwrap();
1408 assert_eq!(reset, 0, "a fresh claim must not be reset");
1409 let status: String = conn
1410 .query_row(
1411 "SELECT status FROM queue WHERE id=?1",
1412 rusqlite::params![id],
1413 |r| r.get(0),
1414 )
1415 .unwrap();
1416 assert_eq!(status, "processing");
1417 let _ = std::fs::remove_file(&path);
1418 }
1419
1420 #[test]
1425 fn enqueue_batch_is_atomic() {
1426 let main = open_test_db();
1427 main.execute(
1428 "INSERT INTO memories (namespace, name, body) VALUES ('global', 'm1', 'b'), ('global', 'm2', 'b'), ('global', 'm3', 'b')",
1429 [],
1430 )
1431 .unwrap();
1432 let (mut queue, path) = open_temp_queue();
1433 let reader = Connection::open(&path).unwrap();
1435
1436 let tx = queue.transaction().unwrap();
1437 let tx_conn: &Connection = &tx;
1438 for key in ["m1", "m2", "m3"] {
1439 enqueue_candidate(tx_conn, &main, "global", key, "memory", "MemoryBindings");
1440 }
1441 let before: i64 = reader
1443 .query_row("SELECT COUNT(*) FROM queue", [], |r| r.get::<_, i64>(0))
1444 .unwrap();
1445 assert_eq!(before, 0, "rows must not be visible before commit");
1446 tx.commit().unwrap();
1447
1448 let after: i64 = reader
1449 .query_row("SELECT COUNT(*) FROM queue", [], |r| r.get::<_, i64>(0))
1450 .unwrap();
1451 assert_eq!(after, 3, "all three rows are visible after commit");
1452 let _ = std::fs::remove_file(&path);
1453 }
1454
1455 #[test]
1458 fn heartbeat_updates_claimed_at() {
1459 let (conn, path) = open_temp_queue();
1460 let id = insert_pending(&conn, "mem-heartbeat");
1461 conn.execute(
1463 "UPDATE queue SET status='processing', claimed_at = CAST(strftime('%s','now') AS INTEGER) - 7200 WHERE id=?1",
1464 rusqlite::params![id],
1465 )
1466 .unwrap();
1467 heartbeat(&conn, id).unwrap();
1468 let claimed_at: Option<i64> = conn
1469 .query_row(
1470 "SELECT claimed_at FROM queue WHERE id=?1",
1471 rusqlite::params![id],
1472 |r| r.get(0),
1473 )
1474 .unwrap();
1475 let claimed_at = claimed_at.expect("claimed_at must be set after heartbeat");
1476 let now: i64 = conn
1477 .query_row("SELECT CAST(strftime('%s','now') AS INTEGER)", [], |r| {
1478 r.get::<_, i64>(0)
1479 })
1480 .unwrap();
1481 assert!(
1482 now - claimed_at <= 5,
1483 "claimed_at must be within 5s of now after heartbeat"
1484 );
1485 let reset = reset_stale_processing_claims(&conn, 1800).unwrap();
1487 assert_eq!(reset, 0, "fresh claim survives the sweep after heartbeat");
1488 let _ = std::fs::remove_file(&path);
1489 }
1490}