use crate::commands::ingest::IngestArgs;
use crate::entity_type::EntityType;
use crate::errors::AppError;
use crate::paths::AppPaths;
use crate::storage::connection::{ensure_db_ready, open_rw};
use crate::storage::entities::{self, NewEntity, NewRelationship};
use crate::storage::memories::{self, NewMemory};
use rusqlite::Connection;
use serde::{Deserialize, Serialize};
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::time::Instant;
#[allow(dead_code)]
const MIN_CLAUDE_VERSION: &str = "2.1.0";
const EXTRACTION_SCHEMA: &str = r#"{
"type": "object",
"properties": {
"name": { "type": "string" },
"description": { "type": "string" },
"entities": {
"type": "array",
"items": {
"type": "object",
"properties": {
"name": { "type": "string" },
"entity_type": {
"type": "string",
"enum": ["project","tool","person","file","concept","incident","decision","organization","location","date"]
}
},
"required": ["name", "entity_type"]
}
},
"relationships": {
"type": "array",
"items": {
"type": "object",
"properties": {
"source": { "type": "string" },
"target": { "type": "string" },
"relation": {
"type": "string",
"enum": ["applies-to","uses","depends-on","causes","fixes","contradicts","supports","follows","related","replaces","tracked-in"]
},
"strength": { "type": "number", "minimum": 0, "maximum": 1 }
},
"required": ["source","target","relation","strength"]
}
}
},
"required": ["name","description","entities","relationships"]
}"#;
const EXTRACTION_PROMPT: &str = "You are a knowledge graph entity extractor. Given a document, extract:\n\
1. A short kebab-case name (max 60 chars) capturing the document's main topic\n\
2. A one-sentence description (10-20 words) summarizing the key insight\n\
3. Domain-specific entities (concepts, tools, people, decisions, projects, files)\n\
4. Typed relationships between entities with strength scores\n\n\
Rules:\n\
- Entity names: lowercase kebab-case, 2+ chars, domain-specific only\n\
- NEVER extract generic terms, stop words, numbers, UUIDs, or single characters\n\
- Relationship types MUST be one of: applies-to, uses, depends-on, causes, fixes, contradicts, supports, follows, related, replaces, tracked-in\n\
- NEVER use 'mentions' as relationship type\n\
- Strength: 0.9 for hard dependencies, 0.7 for design relationships, 0.5 for contextual links, 0.3 for weak references\n\
- Prefer fewer high-quality entities over many low-quality ones\n\
- Description must answer: What is this about and WHY does it matter?";
#[derive(Debug, Deserialize)]
struct ClaudeOutputElement {
r#type: Option<String>,
#[allow(dead_code)]
subtype: Option<String>,
#[serde(default)]
is_error: bool,
structured_output: Option<ExtractionResult>,
result: Option<String>,
total_cost_usd: Option<f64>,
error: Option<String>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ExtractionResult {
pub name: String,
pub description: String,
pub entities: Vec<ExtractedEntity>,
pub relationships: Vec<ExtractedRelationship>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ExtractedEntity {
pub name: String,
pub entity_type: String,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ExtractedRelationship {
pub source: String,
pub target: String,
pub relation: String,
pub strength: f64,
}
#[derive(Debug, Serialize)]
struct PhaseEvent<'a> {
phase: &'a str,
#[serde(skip_serializing_if = "Option::is_none")]
claude_path: Option<&'a str>,
#[serde(skip_serializing_if = "Option::is_none")]
version: Option<&'a str>,
#[serde(skip_serializing_if = "Option::is_none")]
dir: Option<&'a str>,
#[serde(skip_serializing_if = "Option::is_none")]
files_total: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
files_new: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
files_existing: Option<usize>,
}
#[derive(Debug, Serialize)]
struct FileEvent<'a> {
file: &'a str,
name: &'a str,
status: &'a str,
#[serde(skip_serializing_if = "Option::is_none")]
memory_id: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
entities: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
rels: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
cost_usd: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
elapsed_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<&'a str>,
index: usize,
total: usize,
}
#[derive(Debug, Serialize)]
struct Summary {
summary: bool,
files_total: usize,
completed: usize,
failed: usize,
skipped: usize,
entities_total: usize,
rels_total: usize,
cost_usd: f64,
elapsed_ms: u64,
}
pub fn find_claude_binary(explicit: Option<&Path>) -> Result<PathBuf, AppError> {
if let Some(p) = explicit {
if p.exists() {
return Ok(p.to_path_buf());
}
return Err(AppError::Validation(format!(
"Claude Code binary not found at explicit path: {}",
p.display()
)));
}
if let Ok(env_path) = std::env::var("SQLITE_GRAPHRAG_CLAUDE_BINARY") {
let p = PathBuf::from(&env_path);
if p.exists() {
return Ok(p);
}
}
let name = if cfg!(windows) {
"claude.exe"
} else {
"claude"
};
if let Some(path_var) = std::env::var_os("PATH") {
for dir in std::env::split_paths(&path_var) {
let candidate = dir.join(name);
if candidate.exists() {
return Ok(candidate);
}
}
}
Err(AppError::Validation(
"Claude Code binary not found in PATH. Install it from https://docs.anthropic.com/claude-code or specify --claude-binary".to_string(),
))
}
fn validate_claude_version(binary: &Path) -> Result<String, AppError> {
let output = Command::new(binary)
.arg("--version")
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.output()
.map_err(AppError::Io)?;
if !output.status.success() {
return Err(AppError::Validation(
"failed to run 'claude --version'".to_string(),
));
}
let version_str = String::from_utf8(output.stdout)
.map_err(|_| AppError::Validation("claude --version output is not UTF-8".to_string()))?;
let version = version_str.trim().to_string();
Ok(version)
}
fn extract_with_claude(
binary: &Path,
file_content: &[u8],
model: Option<&str>,
timeout_secs: u64,
) -> Result<(ExtractionResult, f64), AppError> {
use wait_timeout::ChildExt;
let mut cmd = Command::new(binary);
cmd.env_clear();
for var in &[
"PATH",
"HOME",
"USER",
"SHELL",
"TERM",
"LANG",
"XDG_CONFIG_HOME",
"XDG_DATA_HOME",
"XDG_RUNTIME_DIR",
"ANTHROPIC_API_KEY",
"CLAUDE_CONFIG_DIR",
"TMPDIR",
"TMP",
"TEMP",
"DYLD_FALLBACK_LIBRARY_PATH",
] {
if let Ok(val) = std::env::var(var) {
cmd.env(var, val);
}
}
cmd.arg("-p")
.arg(EXTRACTION_PROMPT)
.arg("--output-format")
.arg("json")
.arg("--json-schema")
.arg(EXTRACTION_SCHEMA)
.arg("--max-turns")
.arg("3")
.arg("--no-session-persistence");
if std::env::var("ANTHROPIC_API_KEY").is_ok() {
cmd.arg("--bare");
} else {
cmd.arg("--dangerously-skip-permissions");
}
if let Some(m) = model {
cmd.arg("--model").arg(m);
}
cmd.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
let mut child = cmd.spawn().map_err(|e| {
AppError::Io(std::io::Error::new(
e.kind(),
format!("failed to spawn claude: {e}"),
))
})?;
let stdin_data = file_content.to_vec();
let mut child_stdin = child
.stdin
.take()
.ok_or_else(|| AppError::Validation("failed to open claude stdin".into()))?;
let stdin_thread = std::thread::spawn(move || -> Result<(), std::io::Error> {
child_stdin.write_all(&stdin_data)?;
Ok(())
});
let timeout = std::time::Duration::from_secs(timeout_secs);
let status = child.wait_timeout(timeout).map_err(AppError::Io)?;
match status {
Some(exit_status) => {
stdin_thread
.join()
.map_err(|_| AppError::Validation("stdin thread panicked".into()))?
.map_err(AppError::Io)?;
let mut stdout_buf = Vec::new();
let mut stderr_buf = Vec::new();
if let Some(mut out) = child.stdout.take() {
std::io::Read::read_to_end(&mut out, &mut stdout_buf).map_err(AppError::Io)?;
}
if let Some(mut err) = child.stderr.take() {
std::io::Read::read_to_end(&mut err, &mut stderr_buf).map_err(AppError::Io)?;
}
if !exit_status.success() {
let stdout_str = String::from_utf8_lossy(&stdout_buf);
if let Ok(elements) = serde_json::from_str::<Vec<ClaudeOutputElement>>(&stdout_str)
{
if let Some(re) = elements
.iter()
.find(|e| e.r#type.as_deref() == Some("result"))
{
if re.is_error {
let err_msg = re
.error
.as_deref()
.or(re.result.as_deref())
.unwrap_or("unknown error");
if err_msg.contains("rate_limit") || err_msg.contains("overloaded") {
return Err(AppError::Validation(format!(
"RATE_LIMITED: {err_msg}"
)));
}
return Err(AppError::Validation(format!(
"claude -p failed: {err_msg}"
)));
}
}
}
let stderr_str = String::from_utf8_lossy(&stderr_buf);
return Err(AppError::Validation(format!(
"claude -p exited with code {:?}: {}",
exit_status.code(),
stderr_str.trim()
)));
}
let stdout = String::from_utf8(stdout_buf)
.map_err(|_| AppError::Validation("claude -p stdout is not valid UTF-8".into()))?;
parse_claude_output(&stdout)
}
None => {
tracing::warn!(target: "ingest", timeout_secs, "claude -p timed out, killing process");
let _ = child.kill();
let _ = child.wait();
let _ = stdin_thread.join();
Err(AppError::Validation(format!(
"claude -p timed out after {timeout_secs} seconds"
)))
}
}
}
fn parse_claude_output(stdout: &str) -> Result<(ExtractionResult, f64), AppError> {
let elements: Vec<ClaudeOutputElement> = serde_json::from_str(stdout).map_err(|e| {
AppError::Validation(format!("failed to parse claude output as JSON array: {e}"))
})?;
let result_elem = elements
.iter()
.find(|e| e.r#type.as_deref() == Some("result"))
.ok_or_else(|| {
AppError::Validation("claude output missing 'result' element".to_string())
})?;
if result_elem.is_error {
let err_msg = result_elem
.error
.as_deref()
.or(result_elem.result.as_deref())
.unwrap_or("unknown error");
if err_msg.contains("rate_limit") || err_msg.contains("overloaded") {
return Err(AppError::Validation(format!("RATE_LIMITED: {err_msg}")));
}
return Err(AppError::Validation(format!(
"claude extraction failed: {err_msg}"
)));
}
let extraction = result_elem
.structured_output
.clone()
.or_else(|| {
result_elem
.result
.as_ref()
.and_then(|text| serde_json::from_str::<ExtractionResult>(text).ok())
})
.ok_or_else(|| {
AppError::Validation("claude result missing structured_output and result field".into())
})?;
let cost = result_elem.total_cost_usd.unwrap_or(0.0);
Ok((extraction, cost))
}
fn emit_json<T: Serialize>(value: &T) {
if let Ok(json) = serde_json::to_string(value) {
let stdout = std::io::stdout();
let mut lock = stdout.lock();
let _ = writeln!(lock, "{json}");
let _ = lock.flush();
}
}
fn collect_matching_files(
dir: &Path,
pattern: &str,
recursive: bool,
max_files: usize,
) -> Result<Vec<PathBuf>, AppError> {
let mut files = Vec::new();
super::ingest::collect_files(dir, pattern, recursive, &mut files)?;
files.sort();
if files.len() > max_files {
return Err(AppError::Validation(format!(
"found {} files, exceeds --max-files cap of {}",
files.len(),
max_files
)));
}
Ok(files)
}
fn open_queue_db(path: &str) -> Result<Connection, AppError> {
let conn = Connection::open(path)?;
conn.execute_batch(
"CREATE TABLE IF NOT EXISTS queue (
id INTEGER PRIMARY KEY AUTOINCREMENT,
file_path TEXT NOT NULL UNIQUE,
name TEXT,
status TEXT NOT NULL DEFAULT 'pending',
memory_id INTEGER,
entities INTEGER DEFAULT 0,
rels INTEGER DEFAULT 0,
error TEXT,
cost_usd REAL DEFAULT 0.0,
attempt INTEGER DEFAULT 0,
elapsed_ms INTEGER,
created_at TEXT DEFAULT (datetime('now')),
done_at TEXT
);
CREATE INDEX IF NOT EXISTS idx_queue_status ON queue(status);",
)?;
Ok(conn)
}
pub fn run_claude_ingest(args: &IngestArgs) -> Result<(), AppError> {
let started = Instant::now();
if !args.dir.exists() {
return Err(AppError::Validation(format!(
"directory not found: {}",
args.dir.display()
)));
}
let claude_binary = find_claude_binary(args.claude_binary.as_deref())?;
let version = validate_claude_version(&claude_binary)?;
tracing::info!(
target: "ingest",
binary = %claude_binary.display(),
version = %version,
"Claude Code binary validated"
);
emit_json(&PhaseEvent {
phase: "validate",
claude_path: claude_binary.to_str(),
version: Some(&version),
dir: None,
files_total: None,
files_new: None,
files_existing: None,
});
let files = collect_matching_files(&args.dir, &args.pattern, args.recursive, args.max_files)?;
let queue_conn = open_queue_db(&args.queue_db)?;
if args.resume {
let reset = queue_conn
.execute(
"UPDATE queue SET status='pending' WHERE status='processing'",
[],
)
.map_err(|e| AppError::Validation(format!("queue resume failed: {e}")))?;
if reset > 0 {
tracing::info!(target: "ingest", count = reset, "reset stuck processing files to pending");
}
}
if args.retry_failed {
let count = queue_conn
.execute(
"UPDATE queue SET status='pending', attempt=0 WHERE status='failed'",
[],
)
.map_err(|e| AppError::Validation(format!("queue retry-failed reset failed: {e}")))?;
tracing::info!(target: "ingest", count, "retrying failed files");
}
if !args.resume && !args.retry_failed {
queue_conn
.execute("DELETE FROM queue", [])
.map_err(|e| AppError::Validation(format!("queue clear failed: {e}")))?;
}
let mut new_count = 0usize;
let mut existing_count = 0usize;
if !args.retry_failed {
for file in &files {
let file_str = file.to_string_lossy().to_string();
let inserted = queue_conn
.execute(
"INSERT OR IGNORE INTO queue (file_path, status) VALUES (?1, 'pending')",
rusqlite::params![file_str],
)
.map_err(|e| AppError::Validation(format!("queue insert failed: {e}")))?;
if inserted > 0 {
new_count += 1;
} else {
existing_count += 1;
}
}
}
emit_json(&PhaseEvent {
phase: "scan",
claude_path: None,
version: None,
dir: args.dir.to_str(),
files_total: Some(files.len()),
files_new: Some(new_count),
files_existing: Some(existing_count),
});
if args.dry_run {
for (idx, file) in files.iter().enumerate() {
let (name, _truncated, _orig) =
super::ingest::derive_kebab_name(file, args.max_name_length);
emit_json(&FileEvent {
file: &file.to_string_lossy(),
name: &name,
status: "preview",
memory_id: None,
entities: None,
rels: None,
cost_usd: None,
elapsed_ms: None,
error: None,
index: idx,
total: files.len(),
});
}
emit_json(&Summary {
summary: true,
files_total: files.len(),
completed: 0,
failed: 0,
skipped: 0,
entities_total: 0,
rels_total: 0,
cost_usd: 0.0,
elapsed_ms: started.elapsed().as_millis() as u64,
});
if !args.keep_queue {
let _ = std::fs::remove_file(&args.queue_db);
}
return Ok(());
}
let paths = AppPaths::resolve(args.db.as_deref())?;
ensure_db_ready(&paths)?;
let conn = open_rw(&paths.db)?;
let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
let memory_type_str = args.r#type.as_str().to_string();
let mut completed = 0usize;
let mut failed = 0usize;
let skipped = 0usize;
let mut entities_total = 0usize;
let mut rels_total = 0usize;
let mut cost_total = 0.0f64;
let total = files.len();
let mut backoff_secs = args.rate_limit_wait;
loop {
let pending: Option<(i64, String)> = queue_conn
.query_row(
"UPDATE queue SET status='processing', attempt=attempt+1 \
WHERE id = (SELECT id FROM queue WHERE status='pending' ORDER BY id LIMIT 1) \
RETURNING id, file_path",
[],
|row| Ok((row.get(0)?, row.get(1)?)),
)
.ok();
let (queue_id, file_path) = match pending {
Some(p) => p,
None => break,
};
let file_started = Instant::now();
let file_content = match std::fs::read(&file_path) {
Ok(c) => c,
Err(e) => {
let err_msg = format!("IO error: {e}");
let _ = queue_conn.execute(
"UPDATE queue SET status='failed', error=?1, done_at=datetime('now') WHERE id=?2",
rusqlite::params![err_msg, queue_id],
);
let current_index = completed + failed + skipped;
failed += 1;
emit_json(&FileEvent {
file: &file_path,
name: "",
status: "failed",
memory_id: None,
entities: None,
rels: None,
cost_usd: None,
elapsed_ms: Some(file_started.elapsed().as_millis() as u64),
error: Some(&err_msg),
index: current_index,
total,
});
if args.fail_fast {
break;
}
continue;
}
};
let max_extract_attempts: u32 = 2;
let mut extraction_result: Option<(ExtractionResult, f64)> = None;
let mut last_extract_err: Option<String> = None;
for attempt in 1..=max_extract_attempts {
match extract_with_claude(
&claude_binary,
&file_content,
args.claude_model.as_deref(),
args.claude_timeout,
) {
Ok(result) => {
extraction_result = Some(result);
break;
}
Err(ref e) if format!("{e}").contains("RATE_LIMITED") => {
last_extract_err = Some(format!("{e}"));
break;
}
Err(e) => {
let msg = format!("{e}");
if attempt < max_extract_attempts {
tracing::warn!(target: "ingest", attempt, error = %msg, "extraction failed, retrying (cold-start workaround)");
std::thread::sleep(std::time::Duration::from_secs(2));
}
last_extract_err = Some(msg);
}
}
}
if let Some((extraction, cost)) = extraction_result {
backoff_secs = args.rate_limit_wait;
let name = &extraction.name;
let ent_count = extraction.entities.len();
let rel_count = extraction.relationships.len();
let new_entities: Vec<NewEntity> = extraction
.entities
.iter()
.filter_map(|e| match e.entity_type.parse::<EntityType>() {
Ok(et) => Some(NewEntity {
name: e.name.clone(),
entity_type: et,
description: None,
}),
Err(_) => {
tracing::warn!(
target: "ingest",
entity = %e.name,
entity_type = %e.entity_type,
"entity type not recognized, skipping"
);
None
}
})
.collect();
let new_relationships: Vec<NewRelationship> = extraction
.relationships
.iter()
.map(|r| NewRelationship {
source: r.source.clone(),
target: r.target.clone(),
relation: r.relation.clone(),
strength: r.strength,
description: None,
})
.collect();
let body_str = String::from_utf8_lossy(&file_content);
let body_hash = blake3::hash(body_str.as_bytes()).to_hex().to_string();
let new_memory = NewMemory {
name: name.clone(),
namespace: namespace.clone(),
memory_type: memory_type_str.clone(),
description: extraction.description.clone(),
body: body_str.to_string(),
body_hash,
session_id: None,
source: "agent".to_string(),
metadata: serde_json::Value::Object(serde_json::Map::new()),
};
let memory_id = match memories::find_by_name_any_state(&conn, &namespace, name)? {
Some((existing_id, is_deleted)) => {
if is_deleted {
memories::clear_deleted_at(&conn, existing_id)?;
}
let (old_name, old_desc, old_body): (String, String, String) = conn.query_row(
"SELECT name, COALESCE(description,''), COALESCE(body,'') FROM memories WHERE id=?1",
rusqlite::params![existing_id],
|r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)),
)?;
memories::update(&conn, existing_id, &new_memory, None)?;
memories::sync_fts_after_update(
&conn,
existing_id,
&old_name,
&old_desc,
&old_body,
&new_memory.name,
&new_memory.description,
&new_memory.body,
)?;
tracing::info!(target: "ingest", name, memory_id = existing_id, "updated existing memory (force-merge)");
existing_id
}
None => match memories::insert(&conn, &new_memory) {
Ok(id) => id,
Err(e) => {
let err_msg = format!("{e}");
let _ = queue_conn.execute(
"UPDATE queue SET status='failed', error=?1, done_at=datetime('now') WHERE id=?2",
rusqlite::params![err_msg, queue_id],
);
let current_index = completed + failed + skipped;
failed += 1;
emit_json(&FileEvent {
file: &file_path,
name,
status: "failed",
memory_id: None,
entities: None,
rels: None,
cost_usd: Some(cost),
elapsed_ms: Some(file_started.elapsed().as_millis() as u64),
error: Some(&err_msg),
index: current_index,
total,
});
cost_total += cost;
if args.fail_fast {
break;
}
continue;
}
},
};
for ent in &new_entities {
if let Ok(eid) = entities::upsert_entity(&conn, &namespace, ent) {
let _ = entities::link_memory_entity(&conn, memory_id, eid);
}
}
for rel in &new_relationships {
crate::parsers::warn_if_non_canonical(&rel.relation);
let src_id = entities::find_entity_id(&conn, &namespace, &rel.source);
let tgt_id = entities::find_entity_id(&conn, &namespace, &rel.target);
if let (Ok(Some(sid)), Ok(Some(tid))) = (src_id, tgt_id) {
let _ = conn.execute(
"INSERT OR IGNORE INTO relationships (namespace, source_id, target_id, relation, weight) VALUES (?1, ?2, ?3, ?4, ?5)",
rusqlite::params![namespace, sid, tid, rel.relation, rel.strength],
);
}
}
let _ = queue_conn.execute(
"UPDATE queue SET status='done', name=?1, memory_id=?2, entities=?3, rels=?4, cost_usd=?5, elapsed_ms=?6, done_at=datetime('now') WHERE id=?7",
rusqlite::params![
name,
memory_id,
ent_count,
rel_count,
cost,
file_started.elapsed().as_millis() as i64,
queue_id
],
);
let current_index = completed + failed + skipped;
completed += 1;
entities_total += ent_count;
rels_total += rel_count;
cost_total += cost;
emit_json(&FileEvent {
file: &file_path,
name,
status: "done",
memory_id: Some(memory_id),
entities: Some(ent_count),
rels: Some(rel_count),
cost_usd: Some(cost),
elapsed_ms: Some(file_started.elapsed().as_millis() as u64),
error: None,
index: current_index,
total,
});
} else if let Some(ref err_str) = last_extract_err {
if err_str.contains("RATE_LIMITED") {
tracing::warn!(
target: "ingest",
wait_seconds = backoff_secs,
"rate limited, waiting before retry"
);
let _ = queue_conn.execute(
"UPDATE queue SET status='pending' WHERE id=?1",
rusqlite::params![queue_id],
);
std::thread::sleep(std::time::Duration::from_secs(backoff_secs));
backoff_secs = (backoff_secs * 2).min(900);
continue;
} else {
let _ = queue_conn.execute(
"UPDATE queue SET status='failed', error=?1, done_at=datetime('now') WHERE id=?2",
rusqlite::params![err_str, queue_id],
);
let current_index = completed + failed + skipped;
failed += 1;
emit_json(&FileEvent {
file: &file_path,
name: "",
status: "failed",
memory_id: None,
entities: None,
rels: None,
cost_usd: None,
elapsed_ms: Some(file_started.elapsed().as_millis() as u64),
error: Some(err_str),
index: current_index,
total,
});
if args.fail_fast {
break;
}
}
}
if let Some(budget) = args.max_cost_usd {
if cost_total >= budget {
tracing::warn!(
target: "ingest",
spent = cost_total,
budget = budget,
"budget exceeded, stopping"
);
break;
}
}
}
emit_json(&Summary {
summary: true,
files_total: total,
completed,
failed,
skipped,
entities_total,
rels_total,
cost_usd: cost_total,
elapsed_ms: started.elapsed().as_millis() as u64,
});
if !args.keep_queue && failed == 0 {
let _ = std::fs::remove_file(&args.queue_db);
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_extraction_schema_valid_json() {
let _: serde_json::Value =
serde_json::from_str(EXTRACTION_SCHEMA).expect("schema must be valid JSON");
}
#[test]
fn test_parse_claude_output_valid() {
let output = r#"[
{"type":"system","subtype":"init"},
{"type":"assistant"},
{"type":"result","is_error":false,"total_cost_usd":0.02,"structured_output":{"name":"test-doc","description":"A test document","entities":[{"name":"test-entity","entity_type":"concept"}],"relationships":[{"source":"test-entity","target":"test-doc","relation":"applies-to","strength":0.8}]}}
]"#;
let (result, cost) = parse_claude_output(output).expect("parse must succeed");
assert_eq!(result.name, "test-doc");
assert_eq!(result.entities.len(), 1);
assert_eq!(result.relationships.len(), 1);
assert!((cost - 0.02).abs() < f64::EPSILON);
}
#[test]
fn test_parse_claude_output_error() {
let output = r#"[
{"type":"system","subtype":"init"},
{"type":"result","is_error":true,"error":"authentication failed"}
]"#;
let err = parse_claude_output(output).unwrap_err();
assert!(format!("{err}").contains("authentication failed"));
}
#[test]
fn test_parse_claude_output_rate_limit() {
let output = r#"[
{"type":"system","subtype":"init"},
{"type":"result","is_error":true,"error":"rate_limit exceeded"}
]"#;
let err = parse_claude_output(output).unwrap_err();
assert!(format!("{err}").contains("RATE_LIMITED"));
}
#[test]
fn test_parse_claude_output_malformed() {
let output = "not json at all";
assert!(parse_claude_output(output).is_err());
}
#[test]
fn test_find_claude_binary_not_found() {
let original_path = std::env::var_os("PATH");
std::env::set_var("PATH", "/nonexistent");
std::env::remove_var("SQLITE_GRAPHRAG_CLAUDE_BINARY");
let result = find_claude_binary(None);
if let Some(p) = original_path {
std::env::set_var("PATH", p);
}
assert!(result.is_err());
}
#[test]
fn test_parse_claude_output_result_fallback() {
let output = r#"[
{"type":"system","subtype":"init"},
{"type":"result","is_error":false,"total_cost_usd":0.01,"structured_output":null,"result":"{\"name\":\"test-fallback\",\"description\":\"A fallback test\",\"entities\":[{\"name\":\"fb-entity\",\"entity_type\":\"concept\"}],\"relationships\":[]}"}
]"#;
let (result, cost) = parse_claude_output(output).expect("result fallback must work");
assert_eq!(result.name, "test-fallback");
assert_eq!(result.entities.len(), 1);
assert!(result.relationships.is_empty());
assert!((cost - 0.01).abs() < f64::EPSILON);
}
#[test]
fn test_parse_claude_output_error_with_result_field() {
let output = r#"[
{"type":"system","subtype":"init"},
{"type":"result","is_error":true,"result":"Not logged in · Please run /login"}
]"#;
let err = parse_claude_output(output).unwrap_err();
let msg = format!("{err}");
assert!(
msg.contains("Not logged in"),
"expected 'Not logged in' in: {msg}"
);
}
#[test]
fn test_extraction_schema_entity_types_match_enum() {
let schema: serde_json::Value = serde_json::from_str(EXTRACTION_SCHEMA).unwrap();
let types = schema["properties"]["entities"]["items"]["properties"]["entity_type"]["enum"]
.as_array()
.expect("schema must have entity_type enum");
for t in types {
let s = t.as_str().unwrap();
assert!(
s.parse::<EntityType>().is_ok(),
"schema entity_type '{s}' not in EntityType enum"
);
}
}
}