1use super::args::EnrichOperation;
4use crate::errors::AppError;
5use rusqlite::Connection;
6use serde::Serialize;
7
8
9pub(crate) fn open_queue_db<P: AsRef<std::path::Path>>(path: P) -> Result<Connection, AppError> {
26 let conn = Connection::open(path)?;
27 crate::pragmas::apply_sidecar_queue_pragmas(&conn)?;
28 conn.execute_batch(
29 "CREATE TABLE IF NOT EXISTS queue (
30 id INTEGER PRIMARY KEY AUTOINCREMENT,
31 namespace TEXT NOT NULL DEFAULT '',
32 item_key TEXT NOT NULL,
33 item_type TEXT NOT NULL DEFAULT 'memory',
34 status TEXT NOT NULL DEFAULT 'pending',
35 memory_id INTEGER,
36 entity_id INTEGER,
37 entities INTEGER DEFAULT 0,
38 rels INTEGER DEFAULT 0,
39 error TEXT,
40 cost_usd REAL DEFAULT 0.0,
41 attempt INTEGER DEFAULT 0,
42 elapsed_ms INTEGER,
43 created_at TEXT DEFAULT (datetime('now')),
44 done_at TEXT,
45 error_class TEXT,
46 next_retry_at TEXT,
47 operation TEXT,
48 finish_reason TEXT,
49 input_tokens INTEGER,
50 output_tokens INTEGER,
51 claimed_at INTEGER,
52 priority INTEGER NOT NULL DEFAULT 0,
53 UNIQUE (namespace, operation, item_key)
54 );
55 CREATE INDEX IF NOT EXISTS idx_enrich_queue_status ON queue(status);",
56 )?;
57 let mut has_error_class = false;
61 let mut has_next_retry_at = false;
62 let mut has_operation = false;
67 let mut has_finish_reason = false;
73 let mut has_input_tokens = false;
74 let mut has_output_tokens = false;
75 let mut has_claimed_at = false;
80 let mut has_priority = false;
82 let mut has_namespace = false;
84 {
85 let mut stmt = conn.prepare("PRAGMA table_info(queue)")?;
86 let names = stmt.query_map([], |r| r.get::<_, String>(1))?;
87 for name in names {
88 match name?.as_str() {
89 "error_class" => has_error_class = true,
90 "next_retry_at" => has_next_retry_at = true,
91 "operation" => has_operation = true,
92 "finish_reason" => has_finish_reason = true,
93 "input_tokens" => has_input_tokens = true,
94 "output_tokens" => has_output_tokens = true,
95 "claimed_at" => has_claimed_at = true,
96 "priority" => has_priority = true,
97 "namespace" => has_namespace = true,
98 _ => {}
99 }
100 }
101 }
102 if !has_error_class {
103 conn.execute_batch("ALTER TABLE queue ADD COLUMN error_class TEXT")?;
104 }
105 if !has_next_retry_at {
106 conn.execute_batch("ALTER TABLE queue ADD COLUMN next_retry_at TEXT")?;
107 }
108 if !has_operation {
109 conn.execute_batch("ALTER TABLE queue ADD COLUMN operation TEXT")?;
110 }
111 if !has_finish_reason {
112 conn.execute_batch("ALTER TABLE queue ADD COLUMN finish_reason TEXT")?;
113 }
114 if !has_input_tokens {
115 conn.execute_batch("ALTER TABLE queue ADD COLUMN input_tokens INTEGER")?;
116 }
117 if !has_output_tokens {
118 conn.execute_batch("ALTER TABLE queue ADD COLUMN output_tokens INTEGER")?;
119 }
120 if !has_claimed_at {
121 conn.execute_batch("ALTER TABLE queue ADD COLUMN claimed_at INTEGER")?;
122 }
123 if !has_priority {
124 conn.execute_batch(
125 "ALTER TABLE queue ADD COLUMN priority INTEGER NOT NULL DEFAULT 0",
126 )?;
127 }
128 conn.execute(
132 "UPDATE queue SET operation = 'LegacyUnscoped' WHERE operation IS NULL OR operation = ''",
133 [],
134 )?;
135 if !has_namespace {
139 migrate_queue_add_namespace(&conn)?;
140 }
141 conn.execute_batch(
142 "CREATE INDEX IF NOT EXISTS idx_enrich_queue_eligible ON queue(status, next_retry_at);
143 CREATE INDEX IF NOT EXISTS idx_enrich_queue_operation ON queue(operation, status);
144 CREATE INDEX IF NOT EXISTS idx_enrich_queue_memory ON queue(memory_id);
145 CREATE INDEX IF NOT EXISTS idx_enrich_queue_priority ON queue(status, priority DESC, id);
146 CREATE INDEX IF NOT EXISTS idx_enrich_queue_ns ON queue(namespace, operation, status)",
147 )?;
148 Ok(conn)
149}
150
151fn migrate_queue_add_namespace(conn: &Connection) -> Result<(), AppError> {
154 tracing::info!(target: "enrich", "migrating enrich queue to namespace-scoped UNIQUE");
155 conn.execute_batch(
156 "BEGIN;
157 CREATE TABLE queue_v120 (
158 id INTEGER PRIMARY KEY AUTOINCREMENT,
159 namespace TEXT NOT NULL DEFAULT '',
160 item_key TEXT NOT NULL,
161 item_type TEXT NOT NULL DEFAULT 'memory',
162 status TEXT NOT NULL DEFAULT 'pending',
163 memory_id INTEGER,
164 entity_id INTEGER,
165 entities INTEGER DEFAULT 0,
166 rels INTEGER DEFAULT 0,
167 error TEXT,
168 cost_usd REAL DEFAULT 0.0,
169 attempt INTEGER DEFAULT 0,
170 elapsed_ms INTEGER,
171 created_at TEXT DEFAULT (datetime('now')),
172 done_at TEXT,
173 error_class TEXT,
174 next_retry_at TEXT,
175 operation TEXT,
176 finish_reason TEXT,
177 input_tokens INTEGER,
178 output_tokens INTEGER,
179 claimed_at INTEGER,
180 priority INTEGER NOT NULL DEFAULT 0,
181 UNIQUE (namespace, operation, item_key)
182 );
183 INSERT OR IGNORE INTO queue_v120 (
184 id, namespace, item_key, item_type, status, memory_id, entity_id,
185 entities, rels, error, cost_usd, attempt, elapsed_ms, created_at,
186 done_at, error_class, next_retry_at, operation, finish_reason,
187 input_tokens, output_tokens, claimed_at, priority
188 )
189 SELECT
190 id, '', item_key, item_type, status, memory_id, entity_id,
191 entities, rels, error, cost_usd, attempt, elapsed_ms, created_at,
192 done_at, error_class, next_retry_at,
193 COALESCE(NULLIF(operation, ''), 'LegacyUnscoped'),
194 finish_reason, input_tokens, output_tokens, claimed_at,
195 COALESCE(priority, 0)
196 FROM queue;
197 DROP TABLE queue;
198 ALTER TABLE queue_v120 RENAME TO queue;
199 COMMIT;",
200 )
201 .map_err(|e| AppError::Validation(crate::i18n::validation::queue_namespace_migration_failed(&e)))?;
202 Ok(())
203}
204
205pub const PRIORITY_HOT: i64 = 100;
208
209pub(super) fn count_priority_pending(
211 queue_conn: &Connection,
212 operation: &str,
213 min_priority: i64,
214) -> Result<i64, rusqlite::Error> {
215 queue_conn.query_row(
217 "SELECT COUNT(*) FROM queue \
218 WHERE status='pending' \
219 AND (operation = ?1 OR operation IS NULL) \
220 AND COALESCE(priority, 0) >= ?2",
221 rusqlite::params![operation, min_priority],
222 |r| r.get(0),
223 )
224}
225
226pub(super) fn enqueue_candidate(
238 queue_conn: &Connection,
239 main_conn: &Connection,
240 namespace: &str,
241 key: &str,
242 item_type: &str,
243 operation: &str,
244) {
245 let memory_id: Option<i64> = if item_type == "memory" {
248 match main_conn.query_row(
249 "SELECT id FROM memories WHERE namespace=?1 AND name=?2 AND deleted_at IS NULL",
250 rusqlite::params![namespace, key],
251 |r| r.get(0),
252 ) {
253 Ok(id) => Some(id),
254 Err(_) => {
255 tracing::warn!(
256 target: "enrich",
257 namespace,
258 key,
259 "enqueue rejected: memory not found in namespace"
260 );
261 return;
262 }
263 }
264 } else if item_type == "entity" {
265 match main_conn.query_row(
266 "SELECT id FROM entities WHERE namespace=?1 AND name=?2",
267 rusqlite::params![namespace, key],
268 |r| r.get::<_, i64>(0),
269 ) {
270 Ok(_) => None,
271 Err(_) => {
272 tracing::warn!(
273 target: "enrich",
274 namespace,
275 key,
276 "enqueue rejected: entity not found in namespace"
277 );
278 return;
279 }
280 }
281 } else {
282 None
284 };
285 if let Err(e) = queue_conn.execute(
286 "INSERT OR IGNORE INTO queue \
287 (namespace, item_key, item_type, status, operation, memory_id, priority) \
288 VALUES (?1, ?2, ?3, 'pending', ?4, ?5, 0)",
289 rusqlite::params![namespace, key, item_type, operation, memory_id],
290 ) {
291 tracing::warn!(target: "enrich", error = %e, "queue insert failed");
292 }
293}
294
295pub(super) fn enqueue_candidate_with_priority(
297 queue_conn: &Connection,
298 key: &str,
299 item_type: &str,
300 operation: &str,
301 priority: i64,
302) {
303 if let Err(e) = queue_conn.execute(
307 "INSERT OR IGNORE INTO queue \
308 (namespace, item_key, item_type, status, operation, memory_id, priority) \
309 VALUES ('', ?1, ?2, 'pending', ?3, NULL, ?4)",
310 rusqlite::params![key, item_type, operation, priority],
311 ) {
312 tracing::warn!(target: "enrich", error = %e, "priority queue insert failed");
313 } else {
314 let _ = queue_conn.execute(
316 "UPDATE queue SET priority = MAX(COALESCE(priority, 0), ?2), status = CASE \
317 WHEN status IN ('done','skipped','dead') THEN status ELSE 'pending' END \
318 WHERE item_key = ?1 AND COALESCE(priority, 0) < ?2",
319 rusqlite::params![key, priority],
320 );
321 }
322}
323
324pub fn reset_stale_processing_claims(
330 conn: &Connection,
331 max_age_secs: u64,
332) -> Result<usize, AppError> {
333 let reset = conn.execute(
334 "UPDATE queue SET status='pending', claimed_at=NULL \
335 WHERE status='processing' AND claimed_at IS NOT NULL \
336 AND CAST(strftime('%s','now') AS INTEGER) - claimed_at > ?1",
337 rusqlite::params![max_age_secs as i64],
338 )?;
339 Ok(reset)
340}
341
342pub fn heartbeat(conn: &Connection, queue_id: i64) -> Result<(), AppError> {
346 conn.execute(
347 "UPDATE queue SET claimed_at = CAST(strftime('%s','now') AS INTEGER) WHERE id = ?1",
348 rusqlite::params![queue_id],
349 )?;
350 Ok(())
351}
352
353pub(super) fn skipped_item_keys(
361 conn: &Connection,
362 operation: &str,
363) -> Result<std::collections::HashSet<String>, AppError> {
364 let mut stmt = conn.prepare(
365 "SELECT item_key FROM queue WHERE status='skipped' AND (operation = ?1 OR operation IS NULL)",
366 )?;
367 let keys = stmt
368 .query_map(rusqlite::params![operation], |r| r.get::<_, String>(0))?
369 .collect::<Result<std::collections::HashSet<String>, _>>()?;
370 Ok(keys)
371}
372
373pub(super) fn item_type_for(operation: &EnrichOperation) -> &'static str {
377 match operation {
378 EnrichOperation::EntityDescriptions => "entity",
379 EnrichOperation::EntityConnect | EnrichOperation::CrossDomainBridges => "entity_pair",
382 _ => "memory",
383 }
384}
385
386pub(super) fn item_type_for_key(key: &str, default: &'static str) -> &'static str {
396 if key.starts_with("pair:") {
397 "entity_pair"
398 } else if key.starts_with("entity:") {
399 "entity"
400 } else if key.starts_with("chunk:") {
401 "chunk"
402 } else {
403 default
404 }
405}
406
407pub fn cleanup_queue_entry(db_path: &std::path::Path, memory_id: i64, name: &str) {
416 let queue_path = crate::paths::sidecar_path(db_path, ".enrich-queue.sqlite");
417 if !queue_path.exists() {
418 return;
419 }
420 match open_queue_db(&queue_path) {
421 Ok(conn) => {
422 if let Err(e) = conn.execute(
423 "DELETE FROM queue WHERE memory_id = ?1 OR item_key = ?2",
424 rusqlite::params![memory_id, name],
425 ) {
426 tracing::warn!(target: "enrich", error = %e, memory_id, "enrich-queue cleanup failed");
427 }
428 }
429 Err(e) => {
430 tracing::warn!(target: "enrich", error = %e, "enrich-queue cleanup skipped (open failed)");
431 }
432 }
433}
434
435pub(super) fn prune_dead_orphans(
446 queue_conn: &Connection,
447 main_conn: &Connection,
448 operation: &str,
449 namespace: &str,
450) -> Result<i64, AppError> {
451 let dead: Vec<(i64, String)> = {
452 let mut stmt = queue_conn.prepare(
453 "SELECT id, item_key FROM queue \
454 WHERE status='dead' AND item_type='memory' \
455 AND (operation = ?1 OR operation IS NULL) ORDER BY id",
456 )?;
457 let rows = stmt
458 .query_map(rusqlite::params![operation], |r| Ok((r.get(0)?, r.get(1)?)))?
459 .collect::<Result<Vec<_>, _>>()?;
460 rows
461 };
462 let mut pruned = 0_i64;
463 for (id, name) in dead {
464 let exists = main_conn
465 .query_row(
466 "SELECT 1 FROM memories WHERE namespace=?1 AND name=?2 AND deleted_at IS NULL",
467 rusqlite::params![namespace, name],
468 |_| Ok(()),
469 )
470 .is_ok();
471 if !exists {
472 queue_conn.execute("DELETE FROM queue WHERE id=?1", rusqlite::params![id])?;
473 pruned += 1;
474 }
475 }
476 if pruned > 0 {
477 let _ = queue_conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);");
478 }
479 Ok(pruned)
480}
481
482pub(super) fn prune_dead_entity_orphans(
488 queue_conn: &Connection,
489 operation: &str,
490) -> Result<i64, AppError> {
491 let pruned = queue_conn.execute(
492 "DELETE FROM queue \
493 WHERE status='dead' AND item_type='entity' \
494 AND (operation = ?1 OR operation IS NULL)",
495 rusqlite::params![operation],
496 )? as i64;
497 if pruned > 0 {
498 let _ = queue_conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);");
499 }
500 Ok(pruned)
501}
502
503#[derive(Debug, Serialize, schemars::JsonSchema)]
513pub struct EnrichStatus {
514 pub(super) status_report: bool,
515 pub(super) operation: String,
516 pub(super) namespace: String,
517 pub(super) unbound_backlog: usize,
518 pub(super) scan_backlog: i64,
525 #[serde(skip_serializing_if = "Option::is_none")]
527 pub(super) scan_backlog_empty: Option<i64>,
528 #[serde(skip_serializing_if = "Option::is_none")]
530 pub(super) scan_backlog_low_quality: Option<i64>,
531 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
533 pub(super) force_redescribe: bool,
534 #[serde(skip_serializing_if = "Option::is_none")]
537 pub(super) quality_pct: Option<f64>,
538 #[serde(skip_serializing_if = "Option::is_none")]
540 pub(super) quality_sample_n: Option<u32>,
541 #[serde(skip_serializing_if = "Option::is_none")]
543 pub(super) scan_backlog_low_grounding_est: Option<i64>,
544 pub(super) queue_pending: i64,
545 pub(super) queue_processing: i64,
546 pub(super) queue_done: i64,
547 pub(super) queue_failed: i64,
548 pub(super) queue_skipped: i64,
549 pub(super) queue_dead: i64,
550 pub(super) eligible_now: i64,
551 pub(super) waiting: i64,
552 pub(super) state: &'static str,
558 pub(super) waiting_items: Vec<WaitingItem>,
561}
562
563#[derive(Debug, Serialize, schemars::JsonSchema)]
565pub struct WaitingItem {
566 pub(super) item_key: String,
567 pub(super) attempt: i64,
568 pub(super) next_retry_at: Option<String>,
569 pub(super) error_class: Option<String>,
570}
571
572#[derive(Debug, Serialize, schemars::JsonSchema)]
574pub struct DeadItem {
575 pub(super) dead_item: bool,
576 pub(super) item_key: String,
577 pub(super) item_type: String,
578 pub(super) attempt: i64,
579 pub(super) error_class: Option<String>,
580 pub(super) error: Option<String>,
581 pub(super) finish_reason: Option<String>,
586 pub(super) input_tokens: Option<i64>,
588 pub(super) output_tokens: Option<i64>,
590}
591
592#[derive(Debug, Serialize, schemars::JsonSchema)]
594pub struct DeadSummary {
595 pub(super) summary: bool,
596 pub(super) operation: String,
597 pub(super) namespace: String,
598 pub(super) action: &'static str,
600 pub(super) dead_total: i64,
601 pub(super) requeued: i64,
602 pub(super) pruned: i64,
606}
607
608pub(super) use super::queue_ops::*;
610
611#[cfg(test)]
612#[path = "queue_tests_a.rs"]
613mod tests_a;
614#[cfg(test)]
615#[path = "queue_tests_b.rs"]
616mod tests_b;