1#![allow(unused_assignments)] use std::path::PathBuf;
7use std::time::Instant;
8
9use rusqlite::Connection;
10
11use super::args::{EnrichArgs, EnrichMode, EnrichOperation, ReEmbedTarget};
12use super::events::*;
13use super::extraction::find_codex_binary;
14use super::postprocess::take_enrich_backend;
15use super::queue::{
16 count_eligible_pending, enqueue_candidate, item_type_for, item_type_for_key, open_queue_db,
17 reconcile_satisfied_reembed_pending, reopen_force_redescribe_candidates, reset_failed_for_op,
18 reset_processing_for_op, reset_stale_processing_claims, skipped_item_keys,
19};
20use super::scan::{self as scan, count_operation_backlog};
21use super::scheduler;
22use super::DEFAULT_RATE_LIMIT_WAIT;
23use crate::commands::ingest_claude::find_claude_binary;
24use crate::errors::AppError;
25use crate::output::emit_json_line as emit_json;
26use crate::paths::AppPaths;
27use crate::storage::connection::{ensure_db_ready, open_rw};
28
29pub fn run(
31 args: &EnrichArgs,
32 llm_backend: crate::cli::LlmBackendChoice,
33 embedding_backend: crate::cli::EmbeddingBackendChoice,
34) -> Result<(), AppError> {
35 if args.print_schema {
37 return crate::print_schema::emit(crate::print_schema::SchemaId::EnrichStatus);
38 }
39
40 if args.ops_gate {
42 for op in scheduler::gate_ops_order() {
43 let mut gate_args = args.clone();
44 gate_args.operation = Some(op);
45 gate_args.ops_gate = false; run(&gate_args, llm_backend, embedding_backend)?;
47 }
48 return Ok(());
49 }
50
51 validate_mode_conditional_flags_enrich(args)?;
54
55 if args.target != ReEmbedTarget::Memories
58 && !matches!(args.operation(), EnrichOperation::ReEmbed)
59 {
60 let target_label = match args.target {
61 ReEmbedTarget::Memories => "memories",
62 ReEmbedTarget::Entities => "entities",
63 ReEmbedTarget::Chunks => "chunks",
64 ReEmbedTarget::All => "all",
65 };
66 return Err(AppError::Validation(
67 crate::i18n::validation::reembed_target_only(target_label),
68 ));
69 }
70
71 if super::status::try_handle_maintenance(args)? {
72 return Ok(());
73 }
74
75 if args.mode() == EnrichMode::OpenRouter && !args.dry_run {
84 let model = args.openrouter_model.as_deref().ok_or_else(|| {
85 AppError::Validation(crate::i18n::validation::openrouter_model_required())
86 })?;
87 let resolved =
88 crate::config::resolve_api_key("openrouter", args.openrouter_api_key.as_deref())
89 .ok_or_else(|| {
90 AppError::Validation(crate::i18n::validation::openrouter_api_key_not_found())
91 })?;
92 crate::embedder::get_openrouter_chat_client(
93 resolved.value,
94 model,
95 args.openrouter_timeout,
96 )?;
97 }
98
99 let started = Instant::now();
100
101 let paths = AppPaths::resolve(args.db.as_deref())?;
102 ensure_db_ready(&paths)?;
103 let conn = open_rw(&paths.db)?;
104 let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
105
106 let wait_secs = args.wait_job_singleton;
112 let force_flag = args.force_job_singleton;
113 let _singleton = crate::lock::acquire_job_singleton(
114 crate::lock::JobType::Enrich,
115 &namespace,
116 &paths.db,
117 wait_secs,
118 force_flag,
119 )?;
120
121 let provider_binary = if args.dry_run || matches!(args.operation(), EnrichOperation::ReEmbed) {
124 None
125 } else {
126 Some(match args.mode() {
127 EnrichMode::ClaudeCode => {
128 let bin = find_claude_binary(args.claude_binary.as_deref())?;
129 let version = crate::commands::claude_runner::validate_claude_version(&bin)?;
130 tracing::info!(target: "enrich", binary = %bin.display(), version = %version, "Claude Code binary validated");
131 emit_json(&PhaseEvent {
132 phase: "validate",
133 binary_path: bin.to_str(),
134 version: Some(&version),
135 items_total: None,
136 items_pending: None,
137 llm_parallelism: None,
138 });
139 bin
140 }
141 EnrichMode::Codex => {
142 let bin = find_codex_binary(args.codex_binary.as_deref())?;
143 emit_json(&PhaseEvent {
144 phase: "validate",
145 binary_path: bin.to_str(),
146 version: None,
147 items_total: None,
148 items_pending: None,
149 llm_parallelism: None,
150 });
151 bin
152 }
153 EnrichMode::Opencode => {
154 let bin = crate::commands::opencode_runner::find_opencode_binary_with_override(
155 args.opencode_binary.as_deref(),
156 )?;
157 emit_json(&PhaseEvent {
158 phase: "validate",
159 binary_path: bin.to_str(),
160 version: None,
161 items_total: None,
162 items_pending: None,
163 llm_parallelism: None,
164 });
165 bin
166 }
167 EnrichMode::OpenRouter => {
168 emit_json(&PhaseEvent {
173 phase: "validate",
174 binary_path: None,
175 version: None,
176 items_total: None,
177 items_pending: None,
178 llm_parallelism: None,
179 });
180 PathBuf::new()
181 }
182 })
183 };
184
185 if args.max_load_check && !args.dry_run && crate::system_load::is_system_saturated() {
189 let load = crate::system_load::load_average_one();
190 let n = crate::system_load::ncpus();
191 return Err(AppError::Validation(
192 crate::i18n::validation::system_load_exceeded(load, n),
193 ));
194 }
195
196 if args.preflight_check
203 && !args.dry_run
204 && !matches!(args.operation(), EnrichOperation::ReEmbed)
205 {
206 let preflight_result = run_preflight_probe(args);
207 match preflight_result {
208 PreflightOutcome::Healthy => {
209 tracing::info!(target: "enrich", mode = ?args.mode(), "preflight probe healthy");
210 }
211 PreflightOutcome::RateLimited { reason, suggestion } => {
212 if let Some(fallback) = args.fallback_mode.clone() {
213 if fallback != args.mode() {
214 return Err(AppError::Validation(
224 crate::i18n::validation::preflight_rate_limit_fallback(
225 &format!("{:?}", args.mode()),
226 &reason,
227 &format!("{fallback:?}"),
228 ),
229 ));
230 }
231 return Err(AppError::Validation(
232 crate::i18n::validation::preflight_rate_limit_same_mode(
233 &format!("{:?}", args.mode()),
234 &reason,
235 ),
236 ));
237 }
238 return Err(AppError::Validation(
239 crate::i18n::validation::preflight_rate_limit_suggestion(
240 &format!("{:?}", args.mode()),
241 &reason,
242 suggestion,
243 ),
244 ));
245 }
246 PreflightOutcome::Error(e) => {
247 return Err(e);
248 }
249 }
250 }
251
252 let max_runtime_secs = args.max_runtime.unwrap_or(3600);
258 let until_deadline = Instant::now() + std::time::Duration::from_secs(max_runtime_secs);
259 let pair_scan_ops = matches!(
260 args.operation(),
261 EnrichOperation::EntityConnect | EnrichOperation::CrossDomainBridges
262 );
263 const ENTITY_CONNECT_SCAN_SOFT_CEILING_SECS: u64 = 120;
265 let scan_deadline = if pair_scan_ops {
266 let soft =
267 Instant::now() + std::time::Duration::from_secs(ENTITY_CONNECT_SCAN_SOFT_CEILING_SECS);
268 Some(soft.min(until_deadline))
269 } else if args.until_empty || args.max_runtime.is_some() {
270 Some(until_deadline)
271 } else {
272 None
273 };
274
275 let pair_op_cli = enrich_operation_cli_name(&args.operation());
276 let mut backlog_degree0_proxy: Option<i64> = None;
277 if pair_scan_ops {
278 let entities_in_namespace: i64 = conn
279 .query_row(
280 "SELECT COUNT(*) FROM entities WHERE namespace = ?1",
281 rusqlite::params![namespace],
282 |r| r.get(0),
283 )
284 .unwrap_or(0);
285 backlog_degree0_proxy =
287 count_operation_backlog(&conn, &args.operation(), &namespace, args.target).ok();
288 emit_json(&ScanStartEvent {
289 phase: "scan_start",
290 operation: pair_op_cli,
291 entities_in_namespace,
292 backlog_degree0_proxy,
293 pair_algorithm: Some("cooccurrence+hub_island"),
294 limit: args.limit,
295 scan_deadline_secs: scan_deadline
296 .map(|d| d.saturating_duration_since(Instant::now()).as_secs()),
297 });
298 }
299
300 let scan_started = Instant::now();
302 let mut scan_result = scan_operation_with_deadline(&conn, &namespace, args, scan_deadline)?;
303 if matches!(args.operation(), EnrichOperation::BodyEnrich) {
312 let q_path = crate::paths::sidecar_path(&paths.db, ".enrich-queue.sqlite");
313 if let Ok(q) = open_queue_db(&q_path) {
314 if let Ok(vetoed) = skipped_item_keys(&q, &format!("{:?}", args.operation())) {
315 scan_result.retain(|k| !vetoed.contains(k));
316 }
317 }
318 }
319 let total = scan_result.len();
320 let scan_elapsed_ms = scan_started.elapsed().as_millis() as u64;
321
322 emit_json(&PhaseEvent {
323 phase: "scan",
324 binary_path: None,
325 version: None,
326 items_total: Some(total),
327 items_pending: Some(total),
328 llm_parallelism: Some(args.llm_parallelism),
329 });
330 if pair_scan_ops {
331 emit_json(&serde_json::json!({
332 "phase": "scan_meta",
333 "operation": pair_op_cli,
334 "pair_algorithm": "cooccurrence+hub_island",
335 "items_total": total,
336 "pairs_enqueued_this_scan": total,
337 "backlog_degree0_proxy": backlog_degree0_proxy,
338 "scan_elapsed_ms": scan_elapsed_ms,
339 "scan_aborted_reason": serde_json::Value::Null,
340 }));
341 }
342
343 if args.dry_run {
345 if total == 0 {
348 let name_filter = scan::resolve_name_filter(args).unwrap_or_default();
349 if !name_filter.is_empty() {
350 emit_json(&serde_json::json!({
351 "matched": 0,
352 "hint": "no candidates matched --names/--entity-names/--memory-names for this operation; verify name space (entity vs memory) and predicates (e.g. empty description unless --force-redescribe)",
353 "operation": format!("{:?}", args.operation()),
354 "names_requested": name_filter,
355 }));
356 }
357 }
358 for (idx, key) in scan_result.iter().enumerate() {
359 emit_json(&ItemEvent {
360 item: key,
361 status: "preview",
362 memory_id: None,
363 entity_id: None,
364 entities: None,
365 rels: None,
366 chars_before: None,
367 chars_after: None,
368 cost_usd: None,
369 elapsed_ms: None,
370 error: None,
371 index: idx,
372 total,
373 });
374 }
375 emit_json(&EnrichSummary {
376 summary: true,
377 operation: format!("{:?}", args.operation()),
378 items_total: total,
379 completed: 0,
380 failed: 0,
381 skipped: 0,
382 cost_usd: 0.0,
383 elapsed_ms: started.elapsed().as_millis() as u64,
384 backend_invoked: take_enrich_backend(),
385 waiting: 0,
386 dead: 0,
387 budget_exhausted: None,
388 pairs_remaining_estimate: None,
389 yields: None,
390 preempted_for_gate: None,
391 });
392 return Ok(());
393 }
394
395 let queue_path = crate::paths::sidecar_path(&paths.db, ".enrich-queue.sqlite");
401 let mut queue_conn = open_queue_db(&queue_path)?;
402
403 {
408 let stale_reset = reset_stale_processing_claims(&queue_conn, args.stale_claim_secs)?;
409 if stale_reset > 0 {
410 tracing::info!(
411 target: "enrich",
412 count = stale_reset,
413 max_age_secs = args.stale_claim_secs,
414 "reset stale processing claims (older than threshold)"
415 );
416 }
417 }
418
419 let op_label = format!("{:?}", args.operation());
422
423 if args.resume {
425 let reset = reset_processing_for_op(&queue_conn, &op_label, &namespace)
426 .map_err(|e| AppError::Validation(crate::i18n::validation::queue_resume_failed(&e)))?;
427 if reset > 0 {
428 tracing::info!(target: "enrich", count = reset, "reset stuck processing items to pending");
429 }
430 }
431
432 if args.retry_failed {
433 let count = reset_failed_for_op(&queue_conn, &op_label, &namespace).map_err(|e| {
434 AppError::Validation(crate::i18n::validation::queue_retry_failed_reset_failed(&e))
435 })?;
436 tracing::info!(target: "enrich", count, "retrying failed items");
437 }
438 if !args.resume && !args.retry_failed && !args.until_empty {
439 queue_conn
440 .execute(
441 "DELETE FROM queue WHERE operation = ?1 \
442 AND (namespace = ?2 OR namespace = '' OR namespace IS NULL)",
443 rusqlite::params![op_label, namespace],
444 )
445 .map_err(|e| AppError::Validation(crate::i18n::validation::queue_clear_failed(&e)))?;
446 }
447
448 if args.force_redescribe && matches!(args.operation(), EnrichOperation::EntityDescriptions) {
452 let reopened = reopen_force_redescribe_candidates(&queue_conn, &namespace, &scan_result);
453 if reopened > 0 {
454 tracing::info!(
455 target: "enrich",
456 reopened,
457 scan = scan_result.len(),
458 "force-redescribe reopened skipped/done candidates (once per run)"
459 );
460 }
461 }
462
463 if matches!(args.operation(), EnrichOperation::ReEmbed) {
465 match reconcile_satisfied_reembed_pending(&conn, &queue_conn, &namespace) {
466 Ok(n) if n > 0 => {
467 tracing::info!(
468 target: "enrich",
469 reconciled = n,
470 "re-embed reconciled pending rows with live embeddings"
471 );
472 }
473 Ok(_) => {}
474 Err(e) => {
475 tracing::warn!(target: "enrich", error = %e, "re-embed reconcile failed");
476 }
477 }
478 }
479
480 let item_type = item_type_for(&args.operation());
486 {
487 let tx = queue_conn.transaction()?;
488 let tx_conn: &Connection = &tx;
491 for key in scan_result.iter() {
492 let it = item_type_for_key(key, item_type);
496 enqueue_candidate(tx_conn, &conn, &namespace, key, it, &op_label);
497 }
498 tx.commit()?;
499 }
500 {
502 let pending_now = count_eligible_pending(&queue_conn, &op_label, &namespace, "");
503 tracing::info!(
504 target: "enrich",
505 scan = scan_result.len(),
506 pending_after_enqueue = pending_now,
507 "enqueue complete"
508 );
509 }
510
511 let parallelism = super::events::resolve_drain_parallelism(args);
512
513 let mut completed = 0usize;
514 #[allow(unused_assignments)]
515 let mut failed = 0usize;
516 #[allow(unused_assignments)]
517 let mut skipped = 0usize;
518 #[allow(unused_assignments)]
519 let mut cost_total = 0.0f64;
520 #[allow(unused_assignments, unused_variables)]
521 let mut oauth_detected = false;
522 let mut counters = super::drain_parallel::DrainCounters::default();
523 let backoff_secs = DEFAULT_RATE_LIMIT_WAIT;
524 let rate_limit_deadline = std::time::Instant::now() + std::time::Duration::from_secs(3600);
525 let enrich_started = std::time::Instant::now();
526
527 let provider_timeout = match args.mode() {
528 EnrichMode::ClaudeCode => args.claude_timeout,
529 EnrichMode::Codex => args.codex_timeout,
530 EnrichMode::Opencode => args.opencode_timeout,
531 EnrichMode::OpenRouter => args.openrouter_timeout,
532 };
533
534 let provider_model: Option<&str> = match args.mode() {
535 EnrichMode::ClaudeCode => args.claude_model.as_deref(),
536 EnrichMode::Codex => args.codex_model.as_deref(),
537 EnrichMode::Opencode => args.opencode_model.as_deref(),
538 EnrichMode::OpenRouter => args.openrouter_model.as_deref(),
539 };
540
541 let backoff_clause: &str = if args.ignore_backoff {
545 ""
546 } else {
547 "AND (next_retry_at IS NULL OR next_retry_at <= datetime('now'))"
548 };
549
550 emit_json(&ConcurrencyEvent {
553 phase: "concurrency",
554 scan_parallelism: 1,
555 drain_parallelism: parallelism as u32,
556 });
557
558 let mut until_empty_iter: u32 = 0;
566 let yield_every = scheduler::resolve_yield_every_n(args.yield_every_n_items);
567 let mut yield_count: u64 = 0;
568 let mut items_since_yield: usize = 0;
569 let mut preempted_for_gate = false;
571 loop {
574 if args.until_empty {
575 until_empty_iter = until_empty_iter.saturating_add(1);
576 if until_empty_iter > 1 {
577 let mut rescan =
581 scan_operation_with_deadline(&conn, &namespace, args, Some(until_deadline))?;
582 if matches!(args.operation(), EnrichOperation::BodyEnrich) {
587 if let Ok(vetoed) = skipped_item_keys(&queue_conn, &op_label) {
588 rescan.retain(|k| !vetoed.contains(k));
589 }
590 }
591 {
593 let tx = queue_conn.transaction()?;
594 let tx_conn: &Connection = &tx;
595 for key in &rescan {
596 let it = item_type_for_key(key, item_type);
597 enqueue_candidate(tx_conn, &conn, &namespace, key, it, &op_label);
598 }
599 tx.commit()?;
600 }
601 }
602 }
603 let completed_before = completed;
604
605 if parallelism > 1 {
609 super::drain_parallel::drain_parallel(
610 args,
611 &paths,
612 &queue_path,
613 &namespace,
614 provider_binary.as_deref(),
615 provider_model,
616 provider_timeout,
617 &op_label,
618 backoff_clause,
619 parallelism,
620 total,
621 llm_backend,
622 embedding_backend,
623 &mut counters,
624 )?;
625 completed = counters.completed;
626 failed = counters.failed;
627 skipped = counters.skipped;
628 cost_total = counters.cost_total;
629 oauth_detected = counters.oauth_detected;
630 } else {
631 super::drain_serial::drain_serial(
632 args,
633 &paths,
634 &conn,
635 &queue_conn,
636 &namespace,
637 provider_binary.as_deref(),
638 provider_model,
639 provider_timeout,
640 &op_label,
641 backoff_clause,
642 item_type,
643 total,
644 llm_backend,
645 embedding_backend,
646 yield_every,
647 &mut counters,
648 &mut items_since_yield,
649 &mut yield_count,
650 &mut preempted_for_gate,
651 enrich_started,
652 until_deadline,
653 rate_limit_deadline,
654 backoff_secs,
655 )?;
656 completed = counters.completed;
657 failed = counters.failed;
658 skipped = counters.skipped;
659 cost_total = counters.cost_total;
660 oauth_detected = counters.oauth_detected;
661 }
662
663 if !args.until_empty {
664 break;
665 }
666 let eligible_remaining =
668 count_eligible_pending(&queue_conn, &op_label, &namespace, backoff_clause);
669 let progressed = completed > completed_before;
670 if std::time::Instant::now() >= until_deadline {
671 tracing::info!(target: "enrich", "until-empty: max-runtime reached, stopping");
672 break;
673 }
674 if !progressed && eligible_remaining == 0 {
675 tracing::info!(target: "enrich", "until-empty: converged (no eligible items remain)");
676 break;
677 }
678 if eligible_remaining == 0 {
679 std::thread::sleep(std::time::Duration::from_secs(1));
681 }
682 } if crate::shutdown_requested() {
693 let reset = reset_processing_for_op(&queue_conn, &op_label, &namespace).unwrap_or(0);
694 let _ = queue_conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);");
695 tracing::info!(
696 target: "enrich",
697 reset,
698 "graceful shutdown: WAL checkpointed, processing claims reset"
699 );
700 }
701
702 let _ = conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);");
703 let _ = queue_conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);");
704
705 let waiting_final: i64 = queue_conn
709 .query_row(
710 "SELECT COUNT(*) FROM queue WHERE status='pending' \
711 AND (operation = ?1 OR operation IS NULL) \
712 AND next_retry_at IS NOT NULL AND next_retry_at > datetime('now')",
713 rusqlite::params![op_label],
714 |r| r.get(0),
715 )
716 .unwrap_or(0);
717 let dead_final: i64 = queue_conn
718 .query_row(
719 "SELECT COUNT(*) FROM queue WHERE status='dead' \
720 AND (operation = ?1 OR operation IS NULL)",
721 rusqlite::params![op_label],
722 |r| r.get(0),
723 )
724 .unwrap_or(0);
725
726 emit_json(&EnrichSummary {
727 summary: true,
728 operation: format!("{:?}", args.operation()),
729 items_total: total,
730 completed,
731 failed,
732 skipped,
733 cost_usd: cost_total,
734 elapsed_ms: started.elapsed().as_millis() as u64,
735 backend_invoked: take_enrich_backend(),
736 waiting: waiting_final,
737 dead: dead_final,
738 budget_exhausted: if pair_scan_ops && std::time::Instant::now() >= until_deadline {
739 Some(true)
740 } else {
741 None
742 },
743 pairs_remaining_estimate: backlog_degree0_proxy,
744 yields: if yield_count > 0 {
745 Some(yield_count)
746 } else {
747 None
748 },
749 preempted_for_gate: if preempted_for_gate { Some(true) } else { None },
750 });
751
752 if failed == 0 {
753 let dead: i64 = queue_conn
756 .query_row("SELECT COUNT(*) FROM queue WHERE status='dead'", [], |r| {
757 r.get(0)
758 })
759 .unwrap_or(0);
760 let skipped_remaining: i64 = queue_conn
766 .query_row(
767 "SELECT COUNT(*) FROM queue WHERE status='skipped'",
768 [],
769 |r| r.get(0),
770 )
771 .unwrap_or(0);
772 if dead == 0 && skipped_remaining == 0 {
773 let _ = std::fs::remove_file(&queue_path);
774 }
775 }
776
777 Ok(())
778}
779
780