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 enqueue_candidate, item_type_for, item_type_for_key,
17 open_queue_db, reset_stale_processing_claims,
18 skipped_item_keys,
19};
20use super::scan::{
21 self as scan, count_operation_backlog,
22};
23use super::scheduler;
24use super::DEFAULT_RATE_LIMIT_WAIT;
25use crate::commands::ingest_claude::find_claude_binary;
26use crate::errors::AppError;
27use crate::output::emit_json_line as emit_json;
28use crate::paths::AppPaths;
29use crate::storage::connection::{ensure_db_ready, open_rw};
30
31pub fn run(
32 args: &EnrichArgs,
33 llm_backend: crate::cli::LlmBackendChoice,
34 embedding_backend: crate::cli::EmbeddingBackendChoice,
35) -> Result<(), AppError> {
36 if args.ops_gate {
38 for op in scheduler::gate_ops_order() {
39 let mut gate_args = args.clone();
40 gate_args.operation = Some(op);
41 gate_args.ops_gate = false; run(&gate_args, llm_backend, embedding_backend)?;
43 }
44 return Ok(());
45 }
46
47 validate_mode_conditional_flags_enrich(args)?;
50
51 if args.target != ReEmbedTarget::Memories
54 && !matches!(args.operation(), EnrichOperation::ReEmbed)
55 {
56 let target_label = match args.target {
57 ReEmbedTarget::Memories => "memories",
58 ReEmbedTarget::Entities => "entities",
59 ReEmbedTarget::Chunks => "chunks",
60 ReEmbedTarget::All => "all",
61 };
62 return Err(AppError::Validation(format!(
63 "--target {target_label} only applies to --operation re-embed"
64 )));
65 }
66
67 if super::status::try_handle_maintenance(args)? {
68 return Ok(());
69 }
70
71 if args.mode() == EnrichMode::OpenRouter && !args.dry_run {
80 let model = args.openrouter_model.as_deref().ok_or_else(|| {
81 AppError::Validation(
82 "--mode openrouter requires --openrouter-model (no default model is allowed)"
83 .into(),
84 )
85 })?;
86 let resolved =
87 crate::config::resolve_api_key("openrouter", args.openrouter_api_key.as_deref())
88 .ok_or_else(|| {
89 AppError::Validation(
90 "OpenRouter API key not found; store it via \
91 `config add-key --provider openrouter`, or pass --openrouter-api-key \
92 (product env is deprecated)"
93 .into(),
94 )
95 })?;
96 crate::embedder::get_openrouter_chat_client(
97 resolved.value,
98 model,
99 args.openrouter_timeout,
100 )?;
101 }
102
103 let started = Instant::now();
104
105 let paths = AppPaths::resolve(args.db.as_deref())?;
106 ensure_db_ready(&paths)?;
107 let conn = open_rw(&paths.db)?;
108 let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
109
110 let wait_secs = args.wait_job_singleton;
116 let force_flag = args.force_job_singleton;
117 let _singleton = crate::lock::acquire_job_singleton(
118 crate::lock::JobType::Enrich,
119 &namespace,
120 &paths.db,
121 wait_secs,
122 force_flag,
123 )?;
124
125 let provider_binary = if args.dry_run || matches!(args.operation(), EnrichOperation::ReEmbed) {
128 None
129 } else {
130 Some(match args.mode() {
131 EnrichMode::ClaudeCode => {
132 let bin = find_claude_binary(args.claude_binary.as_deref())?;
133 let version = crate::commands::claude_runner::validate_claude_version(&bin)?;
134 tracing::info!(target: "enrich", binary = %bin.display(), version = %version, "Claude Code binary validated");
135 emit_json(&PhaseEvent {
136 phase: "validate",
137 binary_path: bin.to_str(),
138 version: Some(&version),
139 items_total: None,
140 items_pending: None,
141 llm_parallelism: None,
142 });
143 bin
144 }
145 EnrichMode::Codex => {
146 let bin = find_codex_binary(args.codex_binary.as_deref())?;
147 emit_json(&PhaseEvent {
148 phase: "validate",
149 binary_path: bin.to_str(),
150 version: None,
151 items_total: None,
152 items_pending: None,
153 llm_parallelism: None,
154 });
155 bin
156 }
157 EnrichMode::Opencode => {
158 let bin = crate::commands::opencode_runner::find_opencode_binary_with_override(
159 args.opencode_binary.as_deref(),
160 )?;
161 emit_json(&PhaseEvent {
162 phase: "validate",
163 binary_path: bin.to_str(),
164 version: None,
165 items_total: None,
166 items_pending: None,
167 llm_parallelism: None,
168 });
169 bin
170 }
171 EnrichMode::OpenRouter => {
172 emit_json(&PhaseEvent {
177 phase: "validate",
178 binary_path: None,
179 version: None,
180 items_total: None,
181 items_pending: None,
182 llm_parallelism: None,
183 });
184 PathBuf::new()
185 }
186 })
187 };
188
189 if args.max_load_check && !args.dry_run && crate::system_load::is_system_saturated() {
193 let load = crate::system_load::load_average_one();
194 let n = crate::system_load::ncpus();
195 return Err(AppError::Validation(format!(
196 "system load average {load:.2} exceeds 2x ncpus ({n}); \
197 pass --no-max-load-check to override (not recommended)"
198 )));
199 }
200
201 if args.preflight_check
208 && !args.dry_run
209 && !matches!(args.operation(), EnrichOperation::ReEmbed)
210 {
211 let preflight_result = run_preflight_probe(args);
212 match preflight_result {
213 PreflightOutcome::Healthy => {
214 tracing::info!(target: "enrich", mode = ?args.mode(), "preflight probe healthy");
215 }
216 PreflightOutcome::RateLimited { reason, suggestion } => {
217 if let Some(fallback) = args.fallback_mode.clone() {
218 if fallback != args.mode() {
219 return Err(AppError::Validation(format!(
229 "preflight detected rate limit on {mode:?}: {reason}; \
230 re-invoke with `--mode {fallback:?}` to use the fallback provider",
231 mode = args.mode()
232 )));
233 }
234 return Err(AppError::Validation(format!(
235 "preflight detected rate limit on {mode:?}: {reason}; \
236 --fallback-mode matches --mode, no recovery possible",
237 mode = args.mode()
238 )));
239 }
240 return Err(AppError::Validation(format!(
241 "preflight detected rate limit on {mode:?}: {reason}; \
242 {suggestion}; pass --fallback-mode codex to recover",
243 mode = args.mode()
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 if args.resume {
420 let reset = queue_conn
421 .execute(
422 "UPDATE queue SET status='pending' WHERE status='processing'",
423 [],
424 )
425 .map_err(|e| AppError::Validation(format!("queue resume failed: {e}")))?;
426 if reset > 0 {
427 tracing::info!(target: "enrich", count = reset, "reset stuck processing items to pending");
428 }
429 }
430
431 if args.retry_failed {
432 let count = queue_conn
433 .execute(
434 "UPDATE queue SET status='pending', attempt=0 WHERE status='failed'",
435 [],
436 )
437 .map_err(|e| AppError::Validation(format!("queue retry-failed reset failed: {e}")))?;
438 tracing::info!(target: "enrich", count, "retrying failed items");
439 }
440
441 if !args.resume && !args.retry_failed && !args.until_empty {
442 queue_conn
443 .execute("DELETE FROM queue", [])
444 .map_err(|e| AppError::Validation(format!("queue clear failed: {e}")))?;
445 }
446
447 let op_label = format!("{:?}", args.operation());
453 let item_type = item_type_for(&args.operation());
454 {
455 let tx = queue_conn.transaction()?;
456 let tx_conn: &Connection = &tx;
459 for key in scan_result.iter() {
460 let it = item_type_for_key(key, item_type);
464 enqueue_candidate(tx_conn, &conn, &namespace, key, it, &op_label);
465 }
466 tx.commit()?;
467 }
468
469 let parallelism = super::events::resolve_drain_parallelism(args);
470
471 let mut completed = 0usize;
472 #[allow(unused_assignments)]
473 let mut failed = 0usize;
474 #[allow(unused_assignments)]
475 let mut skipped = 0usize;
476 #[allow(unused_assignments)]
477 let mut cost_total = 0.0f64;
478 #[allow(unused_assignments, unused_variables)]
479 let mut oauth_detected = false;
480 let mut counters = super::drain_parallel::DrainCounters::default();
481 let backoff_secs = DEFAULT_RATE_LIMIT_WAIT;
482 let rate_limit_deadline = std::time::Instant::now() + std::time::Duration::from_secs(3600);
483 let enrich_started = std::time::Instant::now();
484
485 let provider_timeout = match args.mode() {
486 EnrichMode::ClaudeCode => args.claude_timeout,
487 EnrichMode::Codex => args.codex_timeout,
488 EnrichMode::Opencode => args.opencode_timeout,
489 EnrichMode::OpenRouter => args.openrouter_timeout,
490 };
491
492 let provider_model: Option<&str> = match args.mode() {
493 EnrichMode::ClaudeCode => args.claude_model.as_deref(),
494 EnrichMode::Codex => args.codex_model.as_deref(),
495 EnrichMode::Opencode => args.opencode_model.as_deref(),
496 EnrichMode::OpenRouter => args.openrouter_model.as_deref(),
497 };
498
499 let backoff_clause: &str = if args.ignore_backoff {
503 ""
504 } else {
505 "AND (next_retry_at IS NULL OR next_retry_at <= datetime('now'))"
506 };
507
508 emit_json(&ConcurrencyEvent {
511 phase: "concurrency",
512 scan_parallelism: 1,
513 drain_parallelism: parallelism as u32,
514 });
515
516 let mut until_empty_iter: u32 = 0;
524 let yield_every = scheduler::resolve_yield_every_n(args.yield_every_n_items);
525 let mut yield_count: u64 = 0;
526 let mut items_since_yield: usize = 0;
527 let mut preempted_for_gate = false;
529 loop {
532 if args.until_empty {
533 until_empty_iter = until_empty_iter.saturating_add(1);
534 if until_empty_iter > 1 {
535 let mut rescan =
539 scan_operation_with_deadline(&conn, &namespace, args, Some(until_deadline))?;
540 if matches!(args.operation(), EnrichOperation::BodyEnrich) {
545 if let Ok(vetoed) = skipped_item_keys(&queue_conn, &op_label) {
546 rescan.retain(|k| !vetoed.contains(k));
547 }
548 }
549 {
551 let tx = queue_conn.transaction()?;
552 let tx_conn: &Connection = &tx;
553 for key in &rescan {
554 let it = item_type_for_key(key, item_type);
555 enqueue_candidate(tx_conn, &conn, &namespace, key, it, &op_label);
556 }
557 tx.commit()?;
558 }
559 }
560 }
561 let completed_before = completed;
562
563 if parallelism > 1 {
567 super::drain_parallel::drain_parallel(
568 args,
569 &paths,
570 &queue_path,
571 &namespace,
572 provider_binary.as_deref(),
573 provider_model,
574 provider_timeout,
575 &op_label,
576 backoff_clause,
577 parallelism,
578 total,
579 llm_backend,
580 embedding_backend,
581 &mut counters,
582 )?;
583 completed = counters.completed;
584 failed = counters.failed;
585 skipped = counters.skipped;
586 cost_total = counters.cost_total;
587 oauth_detected = counters.oauth_detected;
588 } else {
589 super::drain_serial::drain_serial(
590 args,
591 &paths,
592 &conn,
593 &queue_conn,
594 &namespace,
595 provider_binary.as_deref(),
596 provider_model,
597 provider_timeout,
598 &op_label,
599 backoff_clause,
600 item_type,
601 total,
602 llm_backend,
603 embedding_backend,
604 yield_every,
605 &mut counters,
606 &mut items_since_yield,
607 &mut yield_count,
608 &mut preempted_for_gate,
609 enrich_started,
610 until_deadline,
611 rate_limit_deadline,
612 backoff_secs,
613 )?;
614 completed = counters.completed;
615 failed = counters.failed;
616 skipped = counters.skipped;
617 cost_total = counters.cost_total;
618 oauth_detected = counters.oauth_detected;
619 }
620
621 if !args.until_empty {
622 break;
623 }
624 let eligible_remaining: i64 = queue_conn
625 .query_row(
626 &format!("SELECT COUNT(*) FROM queue WHERE status='pending' {backoff_clause}"),
627 [],
628 |r| r.get(0),
629 )
630 .unwrap_or(0);
631 let progressed = completed > completed_before;
632 if std::time::Instant::now() >= until_deadline {
633 tracing::info!(target: "enrich", "until-empty: max-runtime reached, stopping");
634 break;
635 }
636 if !progressed && eligible_remaining == 0 {
637 tracing::info!(target: "enrich", "until-empty: converged (no eligible items remain)");
638 break;
639 }
640 if eligible_remaining == 0 {
641 std::thread::sleep(std::time::Duration::from_secs(1));
643 }
644 } if crate::shutdown_requested() {
654 let reset = queue_conn
655 .execute(
656 "UPDATE queue SET status='pending', claimed_at=NULL WHERE status='processing'",
657 [],
658 )
659 .unwrap_or(0);
660 let _ = queue_conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);");
661 tracing::info!(
662 target: "enrich",
663 reset,
664 "graceful shutdown: WAL checkpointed, processing claims reset"
665 );
666 }
667
668 let _ = conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);");
669 let _ = queue_conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);");
670
671 let waiting_final: i64 = queue_conn
675 .query_row(
676 "SELECT COUNT(*) FROM queue WHERE status='pending' \
677 AND (operation = ?1 OR operation IS NULL) \
678 AND next_retry_at IS NOT NULL AND next_retry_at > datetime('now')",
679 rusqlite::params![op_label],
680 |r| r.get(0),
681 )
682 .unwrap_or(0);
683 let dead_final: i64 = queue_conn
684 .query_row(
685 "SELECT COUNT(*) FROM queue WHERE status='dead' \
686 AND (operation = ?1 OR operation IS NULL)",
687 rusqlite::params![op_label],
688 |r| r.get(0),
689 )
690 .unwrap_or(0);
691
692 emit_json(&EnrichSummary {
693 summary: true,
694 operation: format!("{:?}", args.operation()),
695 items_total: total,
696 completed,
697 failed,
698 skipped,
699 cost_usd: cost_total,
700 elapsed_ms: started.elapsed().as_millis() as u64,
701 backend_invoked: take_enrich_backend(),
702 waiting: waiting_final,
703 dead: dead_final,
704 budget_exhausted: if pair_scan_ops && std::time::Instant::now() >= until_deadline {
705 Some(true)
706 } else {
707 None
708 },
709 pairs_remaining_estimate: backlog_degree0_proxy,
710 yields: if yield_count > 0 { Some(yield_count) } else { None },
711 preempted_for_gate: if preempted_for_gate {
712 Some(true)
713 } else {
714 None
715 },
716 });
717
718 if failed == 0 {
719 let dead: i64 = queue_conn
722 .query_row("SELECT COUNT(*) FROM queue WHERE status='dead'", [], |r| {
723 r.get(0)
724 })
725 .unwrap_or(0);
726 let skipped_remaining: i64 = queue_conn
732 .query_row(
733 "SELECT COUNT(*) FROM queue WHERE status='skipped'",
734 [],
735 |r| r.get(0),
736 )
737 .unwrap_or(0);
738 if dead == 0 && skipped_remaining == 0 {
739 let _ = std::fs::remove_file(&queue_path);
740 }
741 }
742
743 Ok(())
744}
745
746