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 {
206 match operation {
207 EnrichOperation::EntityDescriptions => "entity",
208 _ => "memory",
209 }
210}
211
212pub(super) fn item_type_for_key(key: &str, default: &'static str) -> &'static str {
220 if key.starts_with("entity:") {
221 "entity"
222 } else if key.starts_with("chunk:") {
223 "chunk"
224 } else {
225 default
226 }
227}
228
229pub fn cleanup_queue_entry(db_path: &std::path::Path, memory_id: i64, name: &str) {
238 let queue_path = crate::paths::sidecar_path(db_path, ".enrich-queue.sqlite");
239 if !queue_path.exists() {
240 return;
241 }
242 match open_queue_db(&queue_path) {
243 Ok(conn) => {
244 if let Err(e) = conn.execute(
245 "DELETE FROM queue WHERE memory_id = ?1 OR item_key = ?2",
246 rusqlite::params![memory_id, name],
247 ) {
248 tracing::warn!(target: "enrich", error = %e, memory_id, "enrich-queue cleanup failed");
249 }
250 }
251 Err(e) => {
252 tracing::warn!(target: "enrich", error = %e, "enrich-queue cleanup skipped (open failed)");
253 }
254 }
255}
256
257pub(super) fn prune_dead_orphans(
268 queue_conn: &Connection,
269 main_conn: &Connection,
270 operation: &str,
271 namespace: &str,
272) -> Result<i64, AppError> {
273 let dead: Vec<(i64, String)> = {
274 let mut stmt = queue_conn.prepare(
275 "SELECT id, item_key FROM queue \
276 WHERE status='dead' AND item_type='memory' \
277 AND (operation = ?1 OR operation IS NULL) ORDER BY id",
278 )?;
279 let rows = stmt
280 .query_map(rusqlite::params![operation], |r| Ok((r.get(0)?, r.get(1)?)))?
281 .collect::<Result<Vec<_>, _>>()?;
282 rows
283 };
284 let mut pruned = 0_i64;
285 for (id, name) in dead {
286 let exists = main_conn
287 .query_row(
288 "SELECT 1 FROM memories WHERE namespace=?1 AND name=?2 AND deleted_at IS NULL",
289 rusqlite::params![namespace, name],
290 |_| Ok(()),
291 )
292 .is_ok();
293 if !exists {
294 queue_conn.execute("DELETE FROM queue WHERE id=?1", rusqlite::params![id])?;
295 pruned += 1;
296 }
297 }
298 if pruned > 0 {
299 let _ = queue_conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);");
300 }
301 Ok(pruned)
302}
303
304pub(super) fn prune_dead_entity_orphans(
310 queue_conn: &Connection,
311 operation: &str,
312) -> Result<i64, AppError> {
313 let pruned = queue_conn.execute(
314 "DELETE FROM queue \
315 WHERE status='dead' AND item_type='entity' \
316 AND (operation = ?1 OR operation IS NULL)",
317 rusqlite::params![operation],
318 )? as i64;
319 if pruned > 0 {
320 let _ = queue_conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);");
321 }
322 Ok(pruned)
323}
324
325#[derive(Debug, Serialize, schemars::JsonSchema)]
335pub struct EnrichStatus {
336 pub(super) status_report: bool,
337 pub(super) operation: String,
338 pub(super) namespace: String,
339 pub(super) unbound_backlog: usize,
340 pub(super) scan_backlog: i64,
347 pub(super) queue_pending: i64,
348 pub(super) queue_processing: i64,
349 pub(super) queue_done: i64,
350 pub(super) queue_failed: i64,
351 pub(super) queue_skipped: i64,
352 pub(super) queue_dead: i64,
353 pub(super) eligible_now: i64,
354 pub(super) waiting: i64,
355 pub(super) state: &'static str,
361 pub(super) waiting_items: Vec<WaitingItem>,
364}
365
366#[derive(Debug, Serialize, schemars::JsonSchema)]
368pub struct WaitingItem {
369 pub(super) item_key: String,
370 pub(super) attempt: i64,
371 pub(super) next_retry_at: Option<String>,
372 pub(super) error_class: Option<String>,
373}
374
375#[derive(Debug, Serialize, schemars::JsonSchema)]
377pub struct DeadItem {
378 pub(super) dead_item: bool,
379 pub(super) item_key: String,
380 pub(super) item_type: String,
381 pub(super) attempt: i64,
382 pub(super) error_class: Option<String>,
383 pub(super) error: Option<String>,
384 pub(super) finish_reason: Option<String>,
389 pub(super) input_tokens: Option<i64>,
391 pub(super) output_tokens: Option<i64>,
393}
394
395#[derive(Debug, Serialize, schemars::JsonSchema)]
397pub struct DeadSummary {
398 pub(super) summary: bool,
399 pub(super) operation: String,
400 pub(super) namespace: String,
401 pub(super) action: &'static str,
403 pub(super) dead_total: i64,
404 pub(super) requeued: i64,
405 pub(super) pruned: i64,
409}
410
411pub(super) fn classify_enrich_outcome(e: &AppError) -> crate::retry::AttemptOutcome {
422 use crate::retry::AttemptOutcome;
423 match e {
424 AppError::RateLimited { .. } | AppError::Timeout { .. } | AppError::DbBusy(_) => {
425 AttemptOutcome::Transient
426 }
427 AppError::EntityNotYetMaterialized { .. } => AttemptOutcome::Transient,
435 AppError::ProviderError { .. }
441 | AppError::NotFound(_)
442 | AppError::MemoryNotFound { .. }
443 | AppError::MemoryNotFoundById { .. } => AttemptOutcome::HardFailure,
444 AppError::Database(_) => {
448 if crate::storage::utils::is_sqlite_busy(e) {
449 AttemptOutcome::Transient
450 } else {
451 AttemptOutcome::HardFailure
452 }
453 }
454 AppError::Embedding(_) => AttemptOutcome::Transient,
465 _ => AttemptOutcome::HardFailure,
476 }
477}
478
479pub(super) fn record_item_failure(
492 queue_conn: &rusqlite::Connection,
493 queue_id: i64,
494 attempt: i64,
495 max_attempts: u32,
496 err: &AppError,
497) -> crate::retry::AttemptOutcome {
498 let outcome = classify_enrich_outcome(err);
499 let err_str = format!("{err}");
500 record_item_failure_typed(
501 queue_conn,
502 queue_id,
503 attempt,
504 max_attempts,
505 outcome,
506 &err_str,
507 None,
508 None,
509 None,
510 )
511}
512
513#[allow(clippy::too_many_arguments)]
522pub(super) fn record_item_failure_typed(
523 queue_conn: &rusqlite::Connection,
524 queue_id: i64,
525 attempt: i64,
526 max_attempts: u32,
527 outcome: crate::retry::AttemptOutcome,
528 err_str: &str,
529 finish_reason: Option<&str>,
530 input_tokens: Option<i64>,
531 output_tokens: Option<i64>,
532) -> crate::retry::AttemptOutcome {
533 use crate::retry::AttemptOutcome;
534 let error_class = match outcome {
535 AttemptOutcome::Transient => "transient",
536 AttemptOutcome::HardFailure => "permanent",
537 AttemptOutcome::Success => "success",
538 };
539
540 let terminal = matches!(outcome, AttemptOutcome::HardFailure) || attempt >= max_attempts as i64;
541 if terminal {
542 let _ = queue_conn.execute(
543 "UPDATE queue SET status='dead', error=?1, error_class=?2, done_at=datetime('now'), \
544 finish_reason=?3, input_tokens=?4, output_tokens=?5 WHERE id=?6",
545 rusqlite::params![
546 err_str,
547 error_class,
548 finish_reason,
549 input_tokens,
550 output_tokens,
551 queue_id
552 ],
553 );
554 } else {
555 let delay = crate::retry::compute_delay(
556 &crate::retry::RetryConfig::llm_rate_limit(),
557 attempt.max(0) as u32,
558 );
559 let secs = delay.as_secs().max(1);
560 let modifier = format!("+{secs} seconds");
561 let _ = queue_conn.execute(
562 "UPDATE queue SET status='pending', error=?1, error_class=?2, next_retry_at=datetime('now', ?3), \
563 finish_reason=?4, input_tokens=?5, output_tokens=?6 WHERE id=?7",
564 rusqlite::params![
565 err_str,
566 error_class,
567 modifier,
568 finish_reason,
569 input_tokens,
570 output_tokens,
571 queue_id
572 ],
573 );
574 }
575 outcome
576}
577
578pub(super) enum DequeueOutcome {
586 Claimed((i64, String, String, i64)),
587 Empty,
588}
589
590pub(super) fn dequeue_next_pending(
591 queue_conn: &rusqlite::Connection,
592 backoff_clause: &str,
593) -> Result<DequeueOutcome, AppError> {
594 let dequeue_sql = format!(
595 "UPDATE queue SET status='processing', attempt=attempt+1, \
596 claimed_at=CAST(strftime('%s','now') AS INTEGER) \
597 WHERE id = (SELECT id FROM queue WHERE status='pending' {backoff_clause} \
598 ORDER BY id LIMIT 1) \
599 RETURNING id, item_key, item_type, attempt"
600 );
601 match queue_conn.query_row(&dequeue_sql, [], |row| {
602 Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?))
603 }) {
604 Ok(claimed) => Ok(DequeueOutcome::Claimed(claimed)),
605 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(DequeueOutcome::Empty),
606 Err(e) => Err(AppError::Database(e)),
607 }
608}
609
610#[cfg(test)]
615mod tests {
616 use super::*;
617
618 fn open_test_db() -> Connection {
619 let conn = Connection::open_in_memory().expect("in-memory db");
620 conn.execute_batch(
621 "CREATE TABLE memories (
622 id INTEGER PRIMARY KEY AUTOINCREMENT,
623 namespace TEXT NOT NULL DEFAULT 'global',
624 name TEXT NOT NULL,
625 type TEXT NOT NULL DEFAULT 'note',
626 description TEXT NOT NULL DEFAULT '',
627 body TEXT NOT NULL DEFAULT '',
628 body_hash TEXT NOT NULL DEFAULT '',
629 session_id TEXT,
630 source TEXT NOT NULL DEFAULT 'agent',
631 metadata TEXT NOT NULL DEFAULT '{}',
632 created_at INTEGER NOT NULL DEFAULT (unixepoch()),
633 updated_at INTEGER NOT NULL DEFAULT (unixepoch()),
634 deleted_at INTEGER,
635 UNIQUE(namespace, name)
636 );",
637 )
638 .expect("schema creation must succeed");
639 conn
640 }
641
642 fn open_temp_queue() -> (Connection, String) {
643 let path = format!(
644 "/tmp/test-enrich-dl-{}-{}.sqlite",
645 std::process::id(),
646 fastrand::u64(..)
647 );
648 let conn = open_queue_db(&path).expect("queue db must open");
649 (conn, path)
650 }
651
652 fn insert_pending(conn: &Connection, key: &str) -> i64 {
653 conn.execute(
654 "INSERT INTO queue (item_key, item_type, status) VALUES (?1, 'memory', 'pending')",
655 rusqlite::params![key],
656 )
657 .unwrap();
658 conn.last_insert_rowid()
659 }
660
661 #[test]
662 fn queue_db_schema_creates_correctly() {
663 let tmp_path = format!("/tmp/test-enrich-queue-{}.sqlite", std::process::id());
664 let conn = open_queue_db(&tmp_path).expect("queue db must open");
665 let count: i64 = conn
666 .query_row("SELECT COUNT(*) FROM queue", [], |r| r.get(0))
667 .unwrap();
668 assert_eq!(count, 0);
669 let _ = std::fs::remove_file(&tmp_path);
670 }
671
672 #[test]
673 fn classify_rate_limit_is_transient() {
674 let e = AppError::RateLimited {
675 detail: "429".into(),
676 };
677 assert_eq!(
678 classify_enrich_outcome(&e),
679 crate::retry::AttemptOutcome::Transient
680 );
681 }
682
683 #[test]
684 fn classify_timeout_and_dbbusy_are_transient() {
685 let t = AppError::Timeout {
686 operation: "judge".into(),
687 duration_secs: 30,
688 };
689 let b = AppError::DbBusy("locked".into());
690 assert_eq!(
691 classify_enrich_outcome(&t),
692 crate::retry::AttemptOutcome::Transient
693 );
694 assert_eq!(
695 classify_enrich_outcome(&b),
696 crate::retry::AttemptOutcome::Transient
697 );
698 }
699
700 #[test]
701 fn classify_validation_and_parse_are_hard_failure() {
702 let v = AppError::Validation("failed to parse entities array: bad".into());
703 assert_eq!(
704 classify_enrich_outcome(&v),
705 crate::retry::AttemptOutcome::HardFailure
706 );
707 }
708
709 #[test]
710 fn open_queue_db_alter_is_idempotent() {
711 let path = format!(
712 "/tmp/test-enrich-idem-{}-{}.sqlite",
713 std::process::id(),
714 fastrand::u64(..)
715 );
716 let _ = open_queue_db(&path).expect("first open");
717 let conn = open_queue_db(&path).expect("second open is idempotent");
718 let cols: Vec<String> = {
719 let mut stmt = conn.prepare("PRAGMA table_info(queue)").unwrap();
720 stmt.query_map([], |r| r.get::<_, String>(1))
721 .unwrap()
722 .collect::<Result<Vec<_>, _>>()
723 .unwrap()
724 };
725 assert!(cols.iter().any(|c| c == "error_class"));
726 assert!(cols.iter().any(|c| c == "next_retry_at"));
727 let _ = std::fs::remove_file(&path);
728 }
729
730 #[test]
731 fn record_item_failure_hard_marks_dead() {
732 let (conn, path) = open_temp_queue();
733 let id = insert_pending(&conn, "mem-hard");
734 let outcome = record_item_failure(
735 &conn,
736 id,
737 1,
738 5,
739 &AppError::Validation("invalid body".into()),
740 );
741 assert_eq!(outcome, crate::retry::AttemptOutcome::HardFailure);
742 let status: String = conn
743 .query_row(
744 "SELECT status FROM queue WHERE id=?1",
745 rusqlite::params![id],
746 |r| r.get(0),
747 )
748 .unwrap();
749 assert_eq!(status, "dead");
750 let _ = std::fs::remove_file(&path);
751 }
752
753 #[test]
754 fn record_item_failure_transient_reschedules_pending() {
755 let (conn, path) = open_temp_queue();
756 let id = insert_pending(&conn, "mem-transient");
757 let outcome = record_item_failure(
758 &conn,
759 id,
760 1,
761 5,
762 &AppError::RateLimited {
763 detail: "429".into(),
764 },
765 );
766 assert_eq!(outcome, crate::retry::AttemptOutcome::Transient);
767 let (status, future): (String, i64) = conn
768 .query_row(
769 "SELECT status, (next_retry_at > datetime('now')) FROM queue WHERE id=?1",
770 rusqlite::params![id],
771 |r| Ok((r.get(0)?, r.get(1)?)),
772 )
773 .unwrap();
774 assert_eq!(status, "pending");
775 assert_eq!(future, 1, "next_retry_at must be in the future");
776 let _ = std::fs::remove_file(&path);
777 }
778
779 #[test]
780 fn record_item_failure_transient_at_cap_marks_dead() {
781 let (conn, path) = open_temp_queue();
782 let id = insert_pending(&conn, "mem-cap");
783 let outcome = record_item_failure(
784 &conn,
785 id,
786 5,
787 5,
788 &AppError::RateLimited {
789 detail: "429".into(),
790 },
791 );
792 assert_eq!(outcome, crate::retry::AttemptOutcome::Transient);
793 let status: String = conn
794 .query_row(
795 "SELECT status FROM queue WHERE id=?1",
796 rusqlite::params![id],
797 |r| r.get(0),
798 )
799 .unwrap();
800 assert_eq!(status, "dead");
801 let _ = std::fs::remove_file(&path);
802 }
803
804 #[test]
805 fn dequeue_skips_future_retry_and_dead() {
806 let (conn, path) = open_temp_queue();
807 let eligible = insert_pending(&conn, "mem-eligible");
808 let waiting = insert_pending(&conn, "mem-waiting");
809 conn.execute(
810 "UPDATE queue SET next_retry_at=datetime('now', '+3600 seconds') WHERE id=?1",
811 rusqlite::params![waiting],
812 )
813 .unwrap();
814 let dead = insert_pending(&conn, "mem-dead");
815 conn.execute(
816 "UPDATE queue SET status='dead' WHERE id=?1",
817 rusqlite::params![dead],
818 )
819 .unwrap();
820
821 let claimed: Option<i64> = conn
822 .query_row(
823 "UPDATE queue SET status='processing', attempt=attempt+1 \
824 WHERE id = (SELECT id FROM queue WHERE status='pending' \
825 AND (next_retry_at IS NULL OR next_retry_at <= datetime('now')) \
826 ORDER BY id LIMIT 1) \
827 RETURNING id",
828 [],
829 |r| r.get(0),
830 )
831 .ok();
832 assert_eq!(claimed, Some(eligible));
833
834 let second: Option<i64> = conn
835 .query_row(
836 "UPDATE queue SET status='processing', attempt=attempt+1 \
837 WHERE id = (SELECT id FROM queue WHERE status='pending' \
838 AND (next_retry_at IS NULL OR next_retry_at <= datetime('now')) \
839 ORDER BY id LIMIT 1) \
840 RETURNING id",
841 [],
842 |r| r.get(0),
843 )
844 .ok();
845 assert_eq!(second, None);
846 let _ = std::fs::remove_file(&path);
847 }
848
849 #[test]
850 fn classify_validation_never_infers_transience_from_message() {
851 for msg in [
859 "model 'x' returned non-object JSON after repair (got string)",
860 "model 'x' returned content that could not be parsed even after JSON repair",
861 "model 'x' returned no structured content",
862 "LLM result missing 'description' field",
863 "LLM result missing 'enriched_body' field",
864 ] {
865 assert_eq!(
866 classify_enrich_outcome(&AppError::Validation(msg.into())),
867 crate::retry::AttemptOutcome::HardFailure,
868 "expected hard failure for: {msg}"
869 );
870 }
871 }
872
873 #[test]
874 fn classify_embedding_error_is_transient_floor() {
875 assert_eq!(
876 classify_enrich_outcome(&AppError::Embedding("dimension mismatch".into())),
877 crate::retry::AttemptOutcome::Transient
878 );
879 }
880
881 #[test]
884 fn classify_entity_not_yet_materialized_is_transient() {
885 assert_eq!(
886 classify_enrich_outcome(&AppError::EntityNotYetMaterialized {
887 name: "acme".into(),
888 namespace: "global".into(),
889 }),
890 crate::retry::AttemptOutcome::Transient
891 );
892 }
893
894 #[test]
895 fn classify_memory_absence_stays_hard_failure() {
896 assert_eq!(
897 classify_enrich_outcome(&AppError::MemoryNotFound {
898 name: "mem-x".into(),
899 namespace: "global".into(),
900 }),
901 crate::retry::AttemptOutcome::HardFailure
902 );
903 assert_eq!(
904 classify_enrich_outcome(&AppError::MemoryNotFoundById { id: 42 }),
905 crate::retry::AttemptOutcome::HardFailure
906 );
907 assert_eq!(
908 classify_enrich_outcome(&AppError::NotFound("gone".into())),
909 crate::retry::AttemptOutcome::HardFailure
910 );
911 }
912
913 #[test]
914 fn classify_database_busy_is_transient_non_busy_is_hard() {
915 let busy = AppError::Database(rusqlite::Error::SqliteFailure(
916 rusqlite::ffi::Error::new(rusqlite::ffi::SQLITE_BUSY),
917 Some("database is locked".into()),
918 ));
919 assert_eq!(
920 classify_enrich_outcome(&busy),
921 crate::retry::AttemptOutcome::Transient
922 );
923 let constraint = AppError::Database(rusqlite::Error::SqliteFailure(
924 rusqlite::ffi::Error::new(rusqlite::ffi::SQLITE_CONSTRAINT),
925 Some("UNIQUE constraint failed".into()),
926 ));
927 assert_eq!(
928 classify_enrich_outcome(&constraint),
929 crate::retry::AttemptOutcome::HardFailure
930 );
931 }
932
933 #[test]
934 fn record_item_failure_typed_persists_diagnostics_on_dead_letter() {
935 let (conn, path) = open_temp_queue();
936 let id = insert_pending(&conn, "mem-diag");
937 let outcome = record_item_failure_typed(
938 &conn,
939 id,
940 1,
941 5,
942 crate::retry::AttemptOutcome::HardFailure,
943 "truncated response",
944 Some("length"),
945 Some(120),
946 Some(4096),
947 );
948 assert_eq!(outcome, crate::retry::AttemptOutcome::HardFailure);
949 let (status, finish_reason, input_tokens, output_tokens): (
950 String,
951 Option<String>,
952 Option<i64>,
953 Option<i64>,
954 ) = conn
955 .query_row(
956 "SELECT status, finish_reason, input_tokens, output_tokens FROM queue WHERE id=?1",
957 rusqlite::params![id],
958 |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?)),
959 )
960 .unwrap();
961 assert_eq!(status, "dead");
962 assert_eq!(finish_reason.as_deref(), Some("length"));
963 assert_eq!(input_tokens, Some(120));
964 assert_eq!(output_tokens, Some(4096));
965 let _ = std::fs::remove_file(&path);
966 }
967
968 #[test]
969 fn record_item_failure_typed_reschedules_transient_below_max_attempts() {
970 let (conn, path) = open_temp_queue();
976 let id = insert_pending(&conn, "mem-retry");
977 let outcome = record_item_failure_typed(
978 &conn,
979 id,
980 1,
981 5,
982 crate::retry::AttemptOutcome::Transient,
983 "truncated response",
984 Some("length"),
985 Some(120),
986 Some(64),
987 );
988 assert_eq!(outcome, crate::retry::AttemptOutcome::Transient);
989 let (status, error_class, finish_reason, next_retry_at): (
990 String,
991 String,
992 Option<String>,
993 Option<String>,
994 ) = conn
995 .query_row(
996 "SELECT status, error_class, finish_reason, next_retry_at FROM queue WHERE id=?1",
997 rusqlite::params![id],
998 |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?)),
999 )
1000 .unwrap();
1001 assert_eq!(status, "pending");
1002 assert_eq!(error_class, "transient");
1003 assert_eq!(finish_reason.as_deref(), Some("length"));
1004 assert!(
1005 next_retry_at.is_some(),
1006 "a rescheduled item must carry a next_retry_at"
1007 );
1008 let _ = std::fs::remove_file(&path);
1009 }
1010
1011 #[test]
1021 fn with_busy_retry_bounds_dequeue_under_sustained_contention() {
1022 let (conn, path) = open_temp_queue();
1023 insert_pending(&conn, "mem-busy");
1024 conn.pragma_update(None, "busy_timeout", 0i64)
1025 .expect("busy_timeout override must succeed");
1026
1027 let blocker = Connection::open(&path).expect("blocker connection must open");
1031 blocker
1032 .execute_batch("BEGIN EXCLUSIVE;")
1033 .expect("exclusive lock must be acquired");
1034
1035 let calls = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
1036 let calls_clone = std::sync::Arc::clone(&calls);
1037 let result: Result<DequeueOutcome, AppError> =
1038 crate::storage::utils::with_busy_retry(|| {
1039 calls_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1040 dequeue_next_pending(&conn, "")
1041 });
1042
1043 assert!(
1044 matches!(result, Err(AppError::DbBusy(_))),
1045 "sustained SQLITE_BUSY must convert to DbBusy, not hang or silently report Empty"
1046 );
1047 assert_eq!(
1048 calls.load(std::sync::atomic::Ordering::SeqCst),
1049 crate::constants::MAX_SQLITE_BUSY_RETRIES,
1050 "must attempt exactly MAX_SQLITE_BUSY_RETRIES times, never retry unbounded"
1051 );
1052
1053 blocker
1054 .execute_batch("ROLLBACK;")
1055 .expect("releasing the exclusive lock must succeed");
1056 let _ = std::fs::remove_file(&path);
1057 }
1058
1059 #[test]
1060 fn dequeue_next_pending_distinguishes_empty_from_claimed() {
1061 let (conn, path) = open_temp_queue();
1062 let id = insert_pending(&conn, "mem-dequeue");
1063 let claimed = dequeue_next_pending(&conn, "").expect("dequeue must succeed");
1064 match claimed {
1065 DequeueOutcome::Claimed((claimed_id, key, _, _)) => {
1066 assert_eq!(claimed_id, id);
1067 assert_eq!(key, "mem-dequeue");
1068 }
1069 DequeueOutcome::Empty => panic!("expected a claimed row"),
1070 }
1071 let empty = dequeue_next_pending(&conn, "").expect("dequeue must succeed");
1072 assert!(matches!(empty, DequeueOutcome::Empty));
1073 let _ = std::fs::remove_file(&path);
1074 }
1075
1076 #[test]
1077 fn classify_provider_error_and_not_found_are_hard() {
1078 assert_eq!(
1079 classify_enrich_outcome(&AppError::ProviderError {
1080 code: "400".into(),
1081 message: "context length exceeded".into(),
1082 }),
1083 crate::retry::AttemptOutcome::HardFailure
1084 );
1085 assert_eq!(
1086 classify_enrich_outcome(&AppError::NotFound("memory 'gone' not found".into())),
1087 crate::retry::AttemptOutcome::HardFailure
1088 );
1089 }
1090
1091 #[test]
1092 fn open_queue_db_migrates_operation_column() {
1093 let (conn, path) = open_temp_queue();
1094 drop(conn);
1095 let conn = open_queue_db(&path).expect("second open is idempotent");
1096 let cols: Vec<String> = {
1097 let mut stmt = conn.prepare("PRAGMA table_info(queue)").unwrap();
1098 stmt.query_map([], |r| r.get::<_, String>(1))
1099 .unwrap()
1100 .collect::<Result<Vec<_>, _>>()
1101 .unwrap()
1102 };
1103 assert!(cols.iter().any(|c| c == "operation"));
1104 assert!(cols.iter().any(|c| c == "memory_id"));
1105 let _ = std::fs::remove_file(&path);
1106 }
1107
1108 #[test]
1109 fn enqueue_candidate_tags_operation_and_memory_id() {
1110 let main = open_test_db();
1111 main.execute(
1112 "INSERT INTO memories (namespace, name, body) VALUES ('global', 'mem-x', 'body')",
1113 [],
1114 )
1115 .unwrap();
1116 let mem_id: i64 = main
1117 .query_row("SELECT id FROM memories WHERE name='mem-x'", [], |r| {
1118 r.get(0)
1119 })
1120 .unwrap();
1121 let (queue, path) = open_temp_queue();
1122 enqueue_candidate(&queue, &main, "global", "mem-x", "memory", "MemoryBindings");
1123 let (op, mid): (String, i64) = queue
1124 .query_row(
1125 "SELECT operation, memory_id FROM queue WHERE item_key='mem-x'",
1126 [],
1127 |r| Ok((r.get(0)?, r.get(1)?)),
1128 )
1129 .unwrap();
1130 assert_eq!(op, "MemoryBindings");
1131 assert_eq!(mid, mem_id);
1132 let _ = std::fs::remove_file(&path);
1133 }
1134
1135 #[test]
1136 fn requeue_dead_resurrects_dead_rows() {
1137 let (conn, path) = open_temp_queue();
1138 conn.execute(
1139 "INSERT INTO queue (item_key, item_type, status, operation, attempt, error, error_class, next_retry_at) \
1140 VALUES ('mem-dead', 'memory', 'dead', 'MemoryBindings', 8, 'boom', 'permanent', datetime('now'))",
1141 [],
1142 )
1143 .unwrap();
1144 let n = conn
1145 .execute(
1146 "UPDATE queue SET status='pending', attempt=0, next_retry_at=NULL, \
1147 error=NULL, error_class=NULL \
1148 WHERE status='dead' AND (operation = ?1 OR operation IS NULL)",
1149 rusqlite::params!["MemoryBindings"],
1150 )
1151 .unwrap();
1152 assert_eq!(n, 1);
1153 let (status, attempt, nra): (String, i64, Option<String>) = conn
1154 .query_row(
1155 "SELECT status, attempt, next_retry_at FROM queue WHERE item_key='mem-dead'",
1156 [],
1157 |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)),
1158 )
1159 .unwrap();
1160 assert_eq!(status, "pending");
1161 assert_eq!(attempt, 0);
1162 assert!(nra.is_none());
1163 let _ = std::fs::remove_file(&path);
1164 }
1165
1166 #[test]
1167 fn skipped_item_keys_excludes_only_skipped_for_operation() {
1168 let (conn, path) = open_temp_queue();
1173 conn.execute(
1174 "INSERT INTO queue (item_key, item_type, status, operation) VALUES ('mem-vetoed', 'memory', 'skipped', 'BodyEnrich')",
1175 [],
1176 )
1177 .unwrap();
1178 conn.execute(
1179 "INSERT INTO queue (item_key, item_type, status, operation) VALUES ('mem-pending', 'memory', 'pending', 'BodyEnrich')",
1180 [],
1181 )
1182 .unwrap();
1183 conn.execute(
1184 "INSERT INTO queue (item_key, item_type, status, operation) VALUES ('mem-other-op', 'memory', 'skipped', 'MemoryBindings')",
1185 [],
1186 )
1187 .unwrap();
1188 let keys = skipped_item_keys(&conn, "BodyEnrich").unwrap();
1189 assert!(
1190 keys.contains("mem-vetoed"),
1191 "vetoed BodyEnrich item must be excluded from scan"
1192 );
1193 assert!(
1194 !keys.contains("mem-pending"),
1195 "pending item is still actionable"
1196 );
1197 assert!(
1198 !keys.contains("mem-other-op"),
1199 "skipped item from another operation must not leak"
1200 );
1201 assert_eq!(keys.len(), 1);
1202 let _ = std::fs::remove_file(&path);
1203 }
1204
1205 #[test]
1206 fn cascade_cleanup_delete_targets_memory_id_and_name() {
1207 let (conn, path) = open_temp_queue();
1208 conn.execute(
1209 "INSERT INTO queue (item_key, item_type, status, memory_id) VALUES ('by-id', 'memory', 'done', 42)",
1210 [],
1211 )
1212 .unwrap();
1213 conn.execute(
1214 "INSERT INTO queue (item_key, item_type, status) VALUES ('by-name', 'memory', 'pending')",
1215 [],
1216 )
1217 .unwrap();
1218 let removed = conn
1219 .execute(
1220 "DELETE FROM queue WHERE memory_id = ?1 OR item_key = ?2",
1221 rusqlite::params![42_i64, "by-name"],
1222 )
1223 .unwrap();
1224 assert_eq!(removed, 2);
1225 let remaining: i64 = conn
1226 .query_row("SELECT COUNT(*) FROM queue", [], |r| r.get(0))
1227 .unwrap();
1228 assert_eq!(remaining, 0);
1229 let _ = std::fs::remove_file(&path);
1230 }
1231
1232 #[test]
1233 fn item_type_for_maps_entity_and_memory() {
1234 assert_eq!(
1235 item_type_for(&EnrichOperation::EntityDescriptions),
1236 "entity"
1237 );
1238 assert_eq!(item_type_for(&EnrichOperation::MemoryBindings), "memory");
1239 assert_eq!(item_type_for(&EnrichOperation::AugmentBindings), "memory");
1240 assert_eq!(item_type_for(&EnrichOperation::BodyExtract), "memory");
1241 }
1242
1243 #[test]
1246 fn item_type_for_key_honours_reembed_prefixes() {
1247 assert_eq!(item_type_for_key("plain-memory-name", "memory"), "memory");
1248 assert_eq!(
1249 item_type_for_key("entity:tokio-runtime", "memory"),
1250 "entity"
1251 );
1252 assert_eq!(item_type_for_key("chunk:42", "memory"), "chunk");
1253 assert_eq!(item_type_for_key("some-entity", "entity"), "entity");
1254 }
1255
1256 #[test]
1257 fn prune_dead_orphans_removes_only_orphan_memory_rows() {
1258 let main = open_test_db();
1259 main.execute(
1261 "INSERT INTO memories (namespace, name, body) VALUES ('global', 'alive', 'b')",
1262 [],
1263 )
1264 .unwrap();
1265 let (queue, path) = open_temp_queue();
1266 queue
1268 .execute(
1269 "INSERT INTO queue (item_key, item_type, status, operation, error_class) \
1270 VALUES ('gone', 'memory', 'dead', 'MemoryBindings', 'permanent')",
1271 [],
1272 )
1273 .unwrap();
1274 queue
1276 .execute(
1277 "INSERT INTO queue (item_key, item_type, status, operation, error_class) \
1278 VALUES ('alive', 'memory', 'dead', 'MemoryBindings', 'permanent')",
1279 [],
1280 )
1281 .unwrap();
1282 queue
1284 .execute(
1285 "INSERT INTO queue (item_key, item_type, status, operation) \
1286 VALUES ('some-entity', 'entity', 'dead', 'EntityDescriptions')",
1287 [],
1288 )
1289 .unwrap();
1290
1291 let pruned = prune_dead_orphans(&queue, &main, "MemoryBindings", "global").unwrap();
1292 assert_eq!(pruned, 1, "only the orphan memory row is pruned");
1293
1294 let remaining: Vec<String> = {
1295 let mut stmt = queue
1296 .prepare("SELECT item_key FROM queue ORDER BY item_key")
1297 .unwrap();
1298 stmt.query_map([], |r| r.get::<_, String>(0))
1299 .unwrap()
1300 .collect::<Result<Vec<_>, _>>()
1301 .unwrap()
1302 };
1303 assert_eq!(remaining, vec!["alive", "some-entity"]);
1304 let _ = std::fs::remove_file(&path);
1305 }
1306
1307 #[test]
1308 fn prune_dead_entity_orphans_removes_only_entity_dead_rows() {
1309 let (queue, path) = open_temp_queue();
1310 queue
1312 .execute(
1313 "INSERT INTO queue (item_key, item_type, status, operation, error_class) \
1314 VALUES ('entity:foo', 'entity', 'dead', 'ReEmbed', 'permanent')",
1315 [],
1316 )
1317 .unwrap();
1318 queue
1320 .execute(
1321 "INSERT INTO queue (item_key, item_type, status, operation, error_class) \
1322 VALUES ('mem-dead', 'memory', 'dead', 'MemoryBindings', 'permanent')",
1323 [],
1324 )
1325 .unwrap();
1326 queue
1328 .execute(
1329 "INSERT INTO queue (item_key, item_type, status, operation) \
1330 VALUES ('entity:bar', 'entity', 'pending', 'ReEmbed')",
1331 [],
1332 )
1333 .unwrap();
1334
1335 let pruned = prune_dead_entity_orphans(&queue, "ReEmbed").unwrap();
1336 assert_eq!(pruned, 1, "only the entity dead row is pruned");
1337
1338 let remaining: Vec<String> = {
1339 let mut stmt = queue
1340 .prepare("SELECT item_key FROM queue ORDER BY item_key")
1341 .unwrap();
1342 stmt.query_map([], |r| r.get::<_, String>(0))
1343 .unwrap()
1344 .collect::<Result<Vec<_>, _>>()
1345 .unwrap()
1346 };
1347 assert_eq!(remaining, vec!["entity:bar", "mem-dead"]);
1348 let _ = std::fs::remove_file(&path);
1349 }
1350
1351 #[test]
1354 fn stale_processing_claim_is_reset_after_threshold() {
1355 let (conn, path) = open_temp_queue();
1356 let id = insert_pending(&conn, "mem-stale");
1357 conn.execute(
1359 "UPDATE queue SET status='processing', claimed_at = CAST(strftime('%s','now') AS INTEGER) - 7200 WHERE id=?1",
1360 rusqlite::params![id],
1361 )
1362 .unwrap();
1363 let reset = reset_stale_processing_claims(&conn, 1800).unwrap();
1364 assert_eq!(reset, 1, "a stale claim older than the threshold is reset");
1365 let (status, claimed_at): (String, Option<i64>) = conn
1366 .query_row(
1367 "SELECT status, claimed_at FROM queue WHERE id=?1",
1368 rusqlite::params![id],
1369 |r| Ok((r.get(0)?, r.get(1)?)),
1370 )
1371 .unwrap();
1372 assert_eq!(status, "pending");
1373 assert!(claimed_at.is_none(), "claimed_at must be cleared on reset");
1374 let _ = std::fs::remove_file(&path);
1375 }
1376
1377 #[test]
1380 fn fresh_processing_claim_is_preserved() {
1381 let (conn, path) = open_temp_queue();
1382 let id = insert_pending(&conn, "mem-fresh");
1383 conn.execute(
1385 "UPDATE queue SET status='processing', claimed_at = CAST(strftime('%s','now') AS INTEGER) WHERE id=?1",
1386 rusqlite::params![id],
1387 )
1388 .unwrap();
1389 let reset = reset_stale_processing_claims(&conn, 1800).unwrap();
1390 assert_eq!(reset, 0, "a fresh claim must not be reset");
1391 let status: String = conn
1392 .query_row(
1393 "SELECT status FROM queue WHERE id=?1",
1394 rusqlite::params![id],
1395 |r| r.get(0),
1396 )
1397 .unwrap();
1398 assert_eq!(status, "processing");
1399 let _ = std::fs::remove_file(&path);
1400 }
1401
1402 #[test]
1407 fn enqueue_batch_is_atomic() {
1408 let main = open_test_db();
1409 main.execute(
1410 "INSERT INTO memories (namespace, name, body) VALUES ('global', 'm1', 'b'), ('global', 'm2', 'b'), ('global', 'm3', 'b')",
1411 [],
1412 )
1413 .unwrap();
1414 let (mut queue, path) = open_temp_queue();
1415 let reader = Connection::open(&path).unwrap();
1417
1418 let tx = queue.transaction().unwrap();
1419 let tx_conn: &Connection = &tx;
1420 for key in ["m1", "m2", "m3"] {
1421 enqueue_candidate(tx_conn, &main, "global", key, "memory", "MemoryBindings");
1422 }
1423 let before: i64 = reader
1425 .query_row("SELECT COUNT(*) FROM queue", [], |r| r.get::<_, i64>(0))
1426 .unwrap();
1427 assert_eq!(before, 0, "rows must not be visible before commit");
1428 tx.commit().unwrap();
1429
1430 let after: i64 = reader
1431 .query_row("SELECT COUNT(*) FROM queue", [], |r| r.get::<_, i64>(0))
1432 .unwrap();
1433 assert_eq!(after, 3, "all three rows are visible after commit");
1434 let _ = std::fs::remove_file(&path);
1435 }
1436
1437 #[test]
1440 fn heartbeat_updates_claimed_at() {
1441 let (conn, path) = open_temp_queue();
1442 let id = insert_pending(&conn, "mem-heartbeat");
1443 conn.execute(
1445 "UPDATE queue SET status='processing', claimed_at = CAST(strftime('%s','now') AS INTEGER) - 7200 WHERE id=?1",
1446 rusqlite::params![id],
1447 )
1448 .unwrap();
1449 heartbeat(&conn, id).unwrap();
1450 let claimed_at: Option<i64> = conn
1451 .query_row(
1452 "SELECT claimed_at FROM queue WHERE id=?1",
1453 rusqlite::params![id],
1454 |r| r.get(0),
1455 )
1456 .unwrap();
1457 let claimed_at = claimed_at.expect("claimed_at must be set after heartbeat");
1458 let now: i64 = conn
1459 .query_row("SELECT CAST(strftime('%s','now') AS INTEGER)", [], |r| {
1460 r.get::<_, i64>(0)
1461 })
1462 .unwrap();
1463 assert!(
1464 now - claimed_at <= 5,
1465 "claimed_at must be within 5s of now after heartbeat"
1466 );
1467 let reset = reset_stale_processing_claims(&conn, 1800).unwrap();
1469 assert_eq!(reset, 0, "fresh claim survives the sweep after heartbeat");
1470 let _ = std::fs::remove_file(&path);
1471 }
1472}