Skip to main content

sqlite_graphrag/extract/
llm_embedding.rs

1//! LLM-based embedding backend (v1.0.76 default; reworked in v1.0.79 G42).
2//!
3//! `LlmEmbedding` is the production embedding client. It wraps headless
4//! invocations of `claude code` or `codex` and returns f32 vectors of the
5//! active dimensionality (`crate::constants::embedding_dim()`, default 64).
6//!
7//! v1.0.79 (G42) changes:
8//! - S1: the dimensionality is no longer hardcoded here — the single
9//!   source of truth lives in `crate::constants` and the JSON schemas
10//!   are generated dynamically.
11//! - S2: `embed_batch` embeds N numbered texts per LLM call with the
12//!   `{items:[{i,v}]}` schema, collapsing 39 subprocess spawns into 4-5.
13//! - S4: the codex `--output-schema` file is a `tempfile::NamedTempFile`
14//!   with a randomised name created once per client and shared across
15//!   clones via `Arc` — no per-call write+delete, no PID-path races.
16//! - S5: the claude model honours `SQLITE_GRAPHRAG_CLAUDE_EMBED_MODEL`
17//!   (symmetric to the codex env var). ZERO hardcoded models without
18//!   an env override.
19//! - S6: `CLAUDE_CONFIG_DIR` points at an empty managed directory BY
20//!   DEFAULT, because `--strict-mcp-config`/`--mcp-config '{}'` are
21//!   silently ignored upstream (anthropics/claude-code#10787) and a
22//!   full `~/.claude` costs ~223k cache-creation tokens per call.
23//! - S7: the codex `request_user_input` failure mode maps to an
24//!   actionable error instead of an opaque exit 11.
25//! - BLOCO 4: every subprocess uses `kill_on_drop(true)` plus an
26//!   explicit `tokio::time::timeout`, so cancellation never leaks a
27//!   child and a hung LLM cannot stall the pipeline forever.
28//!
29//! OAuth is the only supported credential path. The constructor rejects
30//! `ANTHROPIC_API_KEY` / `OPENAI_API_KEY` in the environment — see
31//! `v1.0.69 (G31) OAuth-Only Enforcement`.
32
33use crate::errors::AppError;
34use serde::Deserialize;
35use std::process::Stdio;
36use std::sync::Arc;
37use tokio::io::AsyncWriteExt;
38use tokio::process::Command;
39
40/// Default per-LLM-call timeout in seconds. Consistent with the
41/// `--claude-timeout` / `--codex-timeout` defaults used by ingest.
42/// Override via `SQLITE_GRAPHRAG_EMBED_TIMEOUT_SECS`.
43const DEFAULT_EMBED_TIMEOUT_SECS: u64 = 300;
44
45fn embed_timeout() -> std::time::Duration {
46    let secs = std::env::var("SQLITE_GRAPHRAG_EMBED_TIMEOUT_SECS")
47        .ok()
48        .and_then(|v| v.parse::<u64>().ok())
49        .filter(|&n| (10..=3_600).contains(&n))
50        .unwrap_or(DEFAULT_EMBED_TIMEOUT_SECS);
51    std::time::Duration::from_secs(secs)
52}
53
54/// G42/S1: single-vector JSON schema generated from the active dim.
55fn build_single_schema(dim: usize) -> String {
56    format!(
57        r#"{{"type":"object","properties":{{"embedding":{{"type":"array","items":{{"type":"number"}},"minItems":{dim},"maxItems":{dim}}}}},"required":["embedding"],"additionalProperties":false}}"#
58    )
59}
60
61/// G42/S2: batch JSON schema `{items:[{i,v}]}`. The `items` array length
62/// is deliberately unconstrained so ONE schema file serves every batch
63/// size (index coverage is validated in Rust after parsing).
64fn build_batch_schema(dim: usize) -> String {
65    format!(
66        r#"{{"type":"object","properties":{{"items":{{"type":"array","items":{{"type":"object","properties":{{"i":{{"type":"integer"}},"v":{{"type":"array","items":{{"type":"number"}},"minItems":{dim},"maxItems":{dim}}}}},"required":["i","v"],"additionalProperties":false}}}}}},"required":["items"],"additionalProperties":false}}"#
67    )
68}
69
70#[derive(Clone, Debug)]
71pub struct LlmEmbedding {
72    /// Which LLM headless binary to spawn. `claude` or `codex`.
73    flavour: EmbeddingFlavour,
74    /// Cached path to the binary to avoid PATH lookups on every call.
75    binary: std::path::PathBuf,
76    /// Model name. Resolved from env overrides at construction time.
77    model: String,
78    /// G42/S4: lazily-created codex `--output-schema` tempfiles, shared
79    /// across clones. Keyed by dim so an env change between tests cannot
80    /// serve a stale schema.
81    codex_schemas: Arc<parking_lot::Mutex<CodexSchemaFiles>>,
82}
83
84#[derive(Debug, Default)]
85struct CodexSchemaFiles {
86    single: Option<(usize, Arc<tempfile::NamedTempFile>)>,
87    batch: Option<(usize, Arc<tempfile::NamedTempFile>)>,
88}
89
90#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
91pub enum EmbeddingFlavour {
92    Claude,
93    Codex,
94}
95
96impl EmbeddingFlavour {
97    pub fn as_str(self) -> &'static str {
98        match self {
99            Self::Claude => "claude",
100            Self::Codex => "codex",
101        }
102    }
103}
104
105#[derive(Debug, Deserialize)]
106struct EmbeddingResponse {
107    embedding: Vec<f32>,
108}
109
110#[derive(Debug, Deserialize)]
111struct BatchEmbeddingResponse {
112    items: Vec<BatchEmbeddingItem>,
113}
114
115#[derive(Debug, Deserialize)]
116struct BatchEmbeddingItem {
117    i: usize,
118    v: Vec<f32>,
119}
120
121/// Follows symlinks and shell-script shim `exec` targets to find
122/// the real ELF binary. Shim wrappers (like `~/.graphrag-shim/codex`)
123/// can strip hardening flags; bypassing them is a security requirement.
124pub fn resolve_real_binary(path: &std::path::Path) -> std::path::PathBuf {
125    if let Ok(canonical) = std::fs::canonicalize(path) {
126        if is_elf_binary(&canonical) {
127            return canonical;
128        }
129        if let Some(exec_target) = extract_exec_target_from_shim(&canonical) {
130            if exec_target.exists() && is_elf_binary(&exec_target) {
131                return exec_target;
132            }
133        }
134        return canonical;
135    }
136    path.to_path_buf()
137}
138
139fn is_elf_binary(path: &std::path::Path) -> bool {
140    std::fs::read(path)
141        .map(|bytes| bytes.len() >= 4 && bytes[..4] == [0x7f, b'E', b'L', b'F'])
142        .unwrap_or(false)
143}
144
145fn extract_exec_target_from_shim(path: &std::path::Path) -> Option<std::path::PathBuf> {
146    let content = std::fs::read_to_string(path).ok()?;
147    if !content.starts_with("#!") {
148        return None;
149    }
150    for line in content.lines().rev() {
151        let trimmed = line.trim();
152        if trimmed.starts_with("exec ") {
153            let after_exec = trimmed.strip_prefix("exec ")?;
154            let binary = after_exec.split_whitespace().next()?;
155            return Some(std::path::PathBuf::from(binary));
156        }
157    }
158    None
159}
160
161/// G42/S5: claude embedding model with env override, symmetric to the
162/// codex `SQLITE_GRAPHRAG_CODEX_EMBED_MODEL` introduced in v1.0.78.
163fn claude_embed_model() -> String {
164    std::env::var("SQLITE_GRAPHRAG_CLAUDE_EMBED_MODEL")
165        .unwrap_or_else(|_| "claude-sonnet-4-6".to_string())
166}
167
168fn codex_embed_model() -> String {
169    std::env::var("SQLITE_GRAPHRAG_CODEX_EMBED_MODEL").unwrap_or_else(|_| "gpt-5.5".to_string())
170}
171
172impl LlmEmbedding {
173    /// Detects which LLM CLI is available on PATH and returns the
174    /// matching embedding client.
175    ///
176    /// v1.0.76: PREFERS `codex` over `claude` because:
177    /// - Claude Code 2.1+ ships a 180k+ token system context (plugins,
178    ///   skills, agents, MCP) that overflows the 200k context window
179    ///   for even trivial embedding prompts and returns "Prompt is too
180    ///   long". (v1.0.79/S6 mitigates this with an empty
181    ///   `CLAUDE_CONFIG_DIR`, but codex stays the lighter default.)
182    /// - Codex 0.134+ is lightweight (~5k system context) and the
183    ///   `StructuredOutput` tool reliably returns the requested vectors.
184    pub fn detect_available() -> Result<Self, AppError> {
185        Self::oauth_only_enforce()?;
186
187        if let Ok(path) = which::which("codex") {
188            return Ok(Self {
189                flavour: EmbeddingFlavour::Codex,
190                binary: resolve_real_binary(&path),
191                model: codex_embed_model(),
192                codex_schemas: Arc::new(parking_lot::Mutex::new(CodexSchemaFiles::default())),
193            });
194        }
195        if let Ok(path) = which::which("claude") {
196            return Ok(Self {
197                flavour: EmbeddingFlavour::Claude,
198                binary: resolve_real_binary(&path),
199                model: claude_embed_model(),
200                codex_schemas: Arc::new(parking_lot::Mutex::new(CodexSchemaFiles::default())),
201            });
202        }
203        Err(AppError::Embedding(
204            "no LLM CLI found on PATH: install `codex` (0.130+) or `claude` (Claude Code 2.1+)"
205                .to_string(),
206        ))
207    }
208
209    pub fn with_codex() -> Result<Self, AppError> {
210        Self::oauth_only_enforce()?;
211        let path = which::which("codex")
212            .map_err(|_| AppError::Embedding("`codex` not found on PATH".to_string()))?;
213        Ok(Self {
214            flavour: EmbeddingFlavour::Codex,
215            binary: resolve_real_binary(&path),
216            model: codex_embed_model(),
217            codex_schemas: Arc::new(parking_lot::Mutex::new(CodexSchemaFiles::default())),
218        })
219    }
220
221    pub fn with_claude() -> Result<Self, AppError> {
222        Self::oauth_only_enforce()?;
223        let path = which::which("claude")
224            .map_err(|_| AppError::Embedding("`claude` not found on PATH".to_string()))?;
225        Ok(Self {
226            flavour: EmbeddingFlavour::Claude,
227            binary: resolve_real_binary(&path),
228            model: claude_embed_model(),
229            codex_schemas: Arc::new(parking_lot::Mutex::new(CodexSchemaFiles::default())),
230        })
231    }
232
233    /// v1.0.69 (G31): refuse to spawn if an API key is set. The CLI
234    /// must use OAuth. The two API-key env vars are NOT in the
235    /// env-clear whitelist, so a parent process that exports them
236    /// will see this error.
237    fn oauth_only_enforce() -> Result<(), AppError> {
238        if std::env::var("ANTHROPIC_API_KEY").is_ok() {
239            return Err(AppError::Validation(
240                "ANTHROPIC_API_KEY is set; v1.0.76 requires OAuth. \
241                 unset it and use `claude login` instead."
242                    .into(),
243            ));
244        }
245        if std::env::var("OPENAI_API_KEY").is_ok() {
246            return Err(AppError::Validation(
247                "OPENAI_API_KEY is set; v1.0.76 requires OAuth. \
248                 unset it and use `codex login` instead."
249                    .into(),
250            ));
251        }
252        Ok(())
253    }
254
255    /// Embeds a single passage (chunk of a memory body). Returns an
256    /// f32 vector of the active dimensionality.
257    pub fn embed_passage(&self, text: &str) -> Result<Vec<f32>, AppError> {
258        self.invoke_with_prefix(crate::constants::PASSAGE_PREFIX, text)
259    }
260
261    /// Embeds a single query. The LLM uses a different prompt prefix
262    /// to disambiguate query from passage.
263    pub fn embed_query(&self, text: &str) -> Result<Vec<f32>, AppError> {
264        self.invoke_with_prefix(crate::constants::QUERY_PREFIX, text)
265    }
266
267    /// G42/S2: embeds a batch of `(global_index, text)` pairs in ONE
268    /// LLM call. Returns `(global_index, vector)` pairs. Async — this
269    /// is the unit of work scheduled by the bounded fan-out in
270    /// `crate::embedder`.
271    ///
272    /// Cancel safety: the future owns its subprocess via
273    /// `kill_on_drop(true)`, so dropping it (e.g. losing a
274    /// `tokio::select!` race against a cancellation token) kills the
275    /// child and leaks nothing.
276    pub async fn embed_batch_async(
277        &self,
278        prefix: &str,
279        batch: &[(usize, String)],
280    ) -> Result<Vec<(usize, Vec<f32>)>, AppError> {
281        let dim = crate::constants::embedding_dim();
282        if batch.is_empty() {
283            return Ok(Vec::new());
284        }
285        if batch.len() == 1 {
286            let (idx, text) = (&batch[0].0, &batch[0].1);
287            let v = self.invoke_single_async(prefix, text, dim).await?;
288            return Ok(vec![(*idx, v)]);
289        }
290
291        let mut prompt = format!(
292            "Generate {dim}-dimensional semantic embedding vectors for each numbered text below.\n\
293             Return a JSON object with an \"items\" array containing EXACTLY {n} items.\n\
294             Each item has \"i\" (the 1-based index) and \"v\" (the {dim}-float vector, values between -1 and 1).\n\n",
295            n = batch.len()
296        );
297        for (pos, (_, text)) in batch.iter().enumerate() {
298            prompt.push_str(&format!("{}: {prefix}{text}\n", pos + 1));
299        }
300
301        let stdout = match self.flavour {
302            EmbeddingFlavour::Claude => {
303                self.invoke_claude(&prompt, &build_batch_schema(dim))
304                    .await?
305            }
306            EmbeddingFlavour::Codex => {
307                let schema = self.codex_schema_file(dim, true)?;
308                self.invoke_codex(&prompt, schema.path()).await?
309            }
310        };
311
312        let parsed: BatchEmbeddingResponse = parse_llm_json(&stdout).map_err(|e| {
313            AppError::Embedding(format!(
314                "LLM batch embedding response parse failed: {e}; raw={stdout}"
315            ))
316        })?;
317        if parsed.items.len() != batch.len() {
318            return Err(AppError::Embedding(format!(
319                "LLM batch returned {} items, expected {} (G42/S2 coverage check)",
320                parsed.items.len(),
321                batch.len()
322            )));
323        }
324        let mut out: Vec<Option<Vec<f32>>> = vec![None; batch.len()];
325        for item in parsed.items {
326            if item.i == 0 || item.i > batch.len() {
327                return Err(AppError::Embedding(format!(
328                    "LLM batch item index {} out of range 1..={}",
329                    item.i,
330                    batch.len()
331                )));
332            }
333            if item.v.len() != dim {
334                return Err(AppError::Embedding(format!(
335                    "LLM batch item {} returned {} dims, expected {dim}; \
336                     refusing to truncate or pad silently (G42/C5)",
337                    item.i,
338                    item.v.len()
339                )));
340            }
341            out[item.i - 1] = Some(item.v);
342        }
343        let mut result = Vec::with_capacity(batch.len());
344        for (pos, slot) in out.into_iter().enumerate() {
345            let v = slot.ok_or_else(|| {
346                AppError::Embedding(format!(
347                    "LLM batch response is missing item index {} (G42/S2 coverage check)",
348                    pos + 1
349                ))
350            })?;
351            result.push((batch[pos].0, v));
352        }
353        Ok(result)
354    }
355
356    fn invoke_with_prefix(&self, prefix: &str, text: &str) -> Result<Vec<f32>, AppError> {
357        let dim = crate::constants::embedding_dim();
358        let inner = self.invoke_single_async(prefix, text, dim);
359        // v1.0.79 (G42/A2): reuse the process-wide multi-thread runtime
360        // instead of building a current-thread runtime PER CALL. Inside
361        // an existing runtime (tests, async commands) block_in_place
362        // keeps the worker pool healthy.
363        match tokio::runtime::Handle::try_current() {
364            Ok(handle) => tokio::task::block_in_place(|| handle.block_on(inner)),
365            Err(_) => crate::embedder::shared_runtime()?.block_on(inner),
366        }
367    }
368
369    async fn invoke_single_async(
370        &self,
371        prefix: &str,
372        text: &str,
373        dim: usize,
374    ) -> Result<Vec<f32>, AppError> {
375        let prompt = format!("{prefix}{text}");
376        let stdout = match self.flavour {
377            EmbeddingFlavour::Claude => {
378                self.invoke_claude(&prompt, &build_single_schema(dim))
379                    .await?
380            }
381            EmbeddingFlavour::Codex => {
382                let schema = self.codex_schema_file(dim, false)?;
383                self.invoke_codex(&prompt, schema.path()).await?
384            }
385        };
386        let parsed: EmbeddingResponse = parse_llm_json(&stdout).map_err(|e| {
387            AppError::Embedding(format!(
388                "LLM embedding response parse failed: {e}; raw={stdout}"
389            ))
390        })?;
391        if parsed.embedding.len() != dim {
392            return Err(AppError::Embedding(format!(
393                "LLM returned {} dims, expected {dim}; \
394                 refusing to truncate or pad silently (G42/C5)",
395                parsed.embedding.len()
396            )));
397        }
398        Ok(parsed.embedding)
399    }
400
401    /// G42/S4: returns the lazily-created, process-shared codex schema
402    /// tempfile for the requested mode. `NamedTempFile` randomises the
403    /// filename (no PID-based collisions) and removes the file on drop
404    /// of the last `Arc` clone.
405    fn codex_schema_file(
406        &self,
407        dim: usize,
408        batch: bool,
409    ) -> Result<Arc<tempfile::NamedTempFile>, AppError> {
410        let mut guard = self.codex_schemas.lock();
411        let slot = if batch {
412            &mut guard.batch
413        } else {
414            &mut guard.single
415        };
416        if let Some((cached_dim, file)) = slot {
417            if *cached_dim == dim {
418                return Ok(Arc::clone(file));
419            }
420        }
421        let content = if batch {
422            build_batch_schema(dim)
423        } else {
424            build_single_schema(dim)
425        };
426        let file = tempfile::Builder::new()
427            .prefix("sqlite-graphrag-embed-schema-")
428            .suffix(".json")
429            .tempfile()
430            .map_err(|e| AppError::Embedding(format!("schema tempfile create failed: {e}")))?;
431        std::fs::write(file.path(), content)
432            .map_err(|e| AppError::Embedding(format!("schema tempfile write failed: {e}")))?;
433        let file = Arc::new(file);
434        *slot = Some((dim, Arc::clone(&file)));
435        Ok(file)
436    }
437
438    async fn invoke_claude(&self, prompt: &str, schema: &str) -> Result<String, AppError> {
439        // v1.0.69 hardening: --strict-mcp-config --mcp-config '{}' --settings
440        // '{"hooks":{}}' --dangerously-skip-permissions.
441        //
442        // v1.0.76 hardening: Claude Code 2.1+ renamed --output-schema to
443        // --json-schema and accepts the schema as an inline JSON string
444        // (NOT a file path). Also pass --output-format json so the
445        // response is a single JSON object on stdout.
446        //
447        // v1.0.79 (G42/S6): CLAUDE_CONFIG_DIR points at an empty managed
448        // directory BY DEFAULT — the MCP-isolation flags above are
449        // silently ignored upstream (anthropics/claude-code#10787) and a
450        // populated ~/.claude costs ~223k cache-creation tokens per call.
451        let mut cmd = Command::new(&self.binary);
452        cmd.arg("-p")
453            .arg(prompt)
454            .arg("--model")
455            .arg(&self.model)
456            .arg("--json-schema")
457            .arg(schema)
458            .arg("--output-format")
459            .arg("json")
460            .arg("--strict-mcp-config")
461            .arg("--mcp-config")
462            .arg(r#"{"mcpServers":{}}"#)
463            .arg("--settings")
464            .arg(r#"{"hooks":{}}"#)
465            .arg("--dangerously-skip-permissions")
466            .env_clear()
467            .env("PATH", std::env::var("PATH").unwrap_or_default())
468            .env("HOME", std::env::var("HOME").unwrap_or_default())
469            .stdin(Stdio::null())
470            .stdout(Stdio::piped())
471            .stderr(Stdio::piped())
472            // BLOCO 4: cancellation (dropped future) must kill the child.
473            .kill_on_drop(true);
474        if let Some(config_dir) = claude_embedding_config_dir() {
475            cmd.env("CLAUDE_CONFIG_DIR", &config_dir);
476        }
477        let output = tokio::time::timeout(embed_timeout(), cmd.output())
478            .await
479            .map_err(|_| {
480                AppError::Embedding(format!(
481                    "claude embedding call timed out after {}s \
482                     (override via SQLITE_GRAPHRAG_EMBED_TIMEOUT_SECS)",
483                    embed_timeout().as_secs()
484                ))
485            })?
486            .map_err(|e| AppError::Embedding(format!("claude spawn failed: {e}")))?;
487        if !output.status.success() {
488            return Err(AppError::Embedding(format!(
489                "claude exited with {}: stderr={}",
490                output.status,
491                String::from_utf8_lossy(&output.stderr)
492            )));
493        }
494        Ok(String::from_utf8_lossy(&output.stdout).into_owned())
495    }
496
497    async fn invoke_codex(
498        &self,
499        prompt: &str,
500        schema_path: &std::path::Path,
501    ) -> Result<String, AppError> {
502        let mut child = build_codex_embedding_command(&self.binary, &self.model, schema_path)
503            .spawn()
504            .map_err(|e| AppError::Embedding(format!("codex spawn failed: {e}")))?;
505        if let Some(mut stdin) = child.stdin.take() {
506            stdin
507                .write_all(prompt.as_bytes())
508                .await
509                .map_err(|e| AppError::Embedding(format!("codex stdin write failed: {e}")))?;
510        }
511        let output = tokio::time::timeout(embed_timeout(), child.wait_with_output())
512            .await
513            .map_err(|_| {
514                AppError::Embedding(format!(
515                    "codex embedding call timed out after {}s \
516                     (override via SQLITE_GRAPHRAG_EMBED_TIMEOUT_SECS)",
517                    embed_timeout().as_secs()
518                ))
519            })?
520            .map_err(|e| AppError::Embedding(format!("codex wait failed: {e}")))?;
521        if !output.status.success() {
522            let stderr = String::from_utf8_lossy(&output.stderr);
523            // G42/S7: the headless spawn can still hit interactive
524            // prompts on some codex builds; map the failure to an
525            // actionable message instead of an opaque exit.
526            if stderr.contains("request_user_input") {
527                return Err(AppError::Embedding(format!(
528                    "codex requested interactive input in a headless embedding call \
529                     (exit {}). This codex build ignores the non-interactive flags; \
530                     upgrade codex (>= 0.134) or switch the embedding backend to \
531                     claude by removing `codex` from PATH or installing `claude`. \
532                     stderr={stderr}",
533                    output.status
534                )));
535            }
536            return Err(AppError::Embedding(format!(
537                "codex exited with {}: stderr={stderr}",
538                output.status
539            )));
540        }
541        Ok(String::from_utf8_lossy(&output.stdout).into_owned())
542    }
543}
544
545/// G42/S6: resolves the empty `CLAUDE_CONFIG_DIR` used for embedding
546/// subprocesses.
547///
548/// - `SQLITE_GRAPHRAG_CLAUDE_EMPTY_CONFIG_DIR` is honoured when set and
549///   pointing at a directory (same contract as G28-A in claude_runner);
550/// - otherwise a managed directory is created at
551///   `~/.local/state/sqlite-graphrag/claude-empty-config` (mode 0700).
552///   If `~/.claude/.credentials.json` exists (Linux OAuth storage) it is
553///   copied in so authentication still works; on macOS credentials live
554///   in the Keychain and the empty dir is sufficient.
555///
556/// Returns `None` only when HOME is unset AND no override is given —
557/// in that case the subprocess falls back to claude's own default.
558fn claude_embedding_config_dir() -> Option<std::path::PathBuf> {
559    if let Ok(dir) = std::env::var("SQLITE_GRAPHRAG_CLAUDE_EMPTY_CONFIG_DIR") {
560        let path = std::path::PathBuf::from(dir);
561        if path.is_dir() {
562            return Some(path);
563        }
564        tracing::warn!(
565            target: "embedding",
566            path = %path.display(),
567            "SQLITE_GRAPHRAG_CLAUDE_EMPTY_CONFIG_DIR is set but not a directory; \
568             falling back to the managed empty config dir"
569        );
570    }
571    let home = std::env::var("HOME").ok()?;
572    let dir = std::path::Path::new(&home)
573        .join(".local/state/sqlite-graphrag")
574        .join("claude-empty-config");
575    if std::fs::create_dir_all(&dir).is_err() {
576        return None;
577    }
578    #[cfg(unix)]
579    {
580        use std::os::unix::fs::PermissionsExt;
581        let _ = std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o700));
582    }
583    // Linux stores OAuth credentials on disk; copy them so the isolated
584    // config dir still authenticates. Best-effort: macOS uses Keychain.
585    let creds = std::path::Path::new(&home).join(".claude/.credentials.json");
586    if creds.exists() {
587        let target = dir.join(".credentials.json");
588        if !target.exists() {
589            let _ = std::fs::copy(&creds, &target);
590        }
591    }
592    Some(dir)
593}
594
595fn build_codex_embedding_command(
596    binary: &std::path::Path,
597    model: &str,
598    schema_path: &std::path::Path,
599) -> Command {
600    let mut cmd = Command::new(binary);
601    // v1.0.77: `-c` TOML overrides bypass the codex exec --sandbox propagation
602    // bug (openai/codex#18113). CLI flags alone are insufficient — the exec
603    // subcommand may not inherit --sandbox from the parent codex command.
604    cmd.arg("exec")
605        .arg("-c")
606        .arg("sandbox_mode='read-only'")
607        .arg("-c")
608        .arg("approval_policy='never'")
609        .arg("--json")
610        .arg("--output-schema")
611        .arg(schema_path)
612        .arg("--ephemeral")
613        .arg("--skip-git-repo-check")
614        .arg("--sandbox")
615        .arg("read-only")
616        .arg("--ignore-user-config")
617        .arg("--ignore-rules");
618    if crate::extract::codex_compat::codex_supports_ask_for_approval() {
619        cmd.arg("--ask-for-approval").arg("never");
620    }
621    // v1.0.77: isolate codex from user config by pointing CODEX_HOME at a
622    // minimal directory containing only auth.json (OAuth credentials).
623    let codex_home = prepare_isolated_codex_home();
624    cmd.arg("--model")
625        .arg(model)
626        .arg("-")
627        .env_clear()
628        .env("PATH", std::env::var("PATH").unwrap_or_default())
629        .env("HOME", std::env::var("HOME").unwrap_or_default());
630    if let Some(ref ch) = codex_home {
631        cmd.env("CODEX_HOME", ch);
632    }
633    cmd.stdin(Stdio::piped())
634        .stdout(Stdio::piped())
635        .stderr(Stdio::piped())
636        // BLOCO 4: cancellation (dropped future) must kill the child.
637        .kill_on_drop(true);
638    cmd
639}
640
641fn prepare_isolated_codex_home() -> Option<std::path::PathBuf> {
642    let home = std::env::var("HOME").ok()?;
643    let real_auth = std::path::Path::new(&home).join(".codex/auth.json");
644    if !real_auth.exists() {
645        return None;
646    }
647    let base = std::path::Path::new(&home).join(".local/share/sqlite-graphrag");
648    let isolated = base.join(format!("codex-home-{}", std::process::id()));
649    let _ = std::fs::create_dir_all(&isolated);
650    let target = isolated.join("auth.json");
651    if !target.exists() {
652        let _ = std::fs::copy(&real_auth, &target);
653    }
654    Some(isolated)
655}
656
657/// Parse an LLM JSON response of type `T`. The two backends emit
658/// different shapes:
659/// - Claude (with `--output-format json`): single JSON object on stdout.
660/// - Codex (with `--json`): JSONL stream with one event per line; the
661///   `agent_message` event's `text` field is the JSON payload.
662///
663/// This helper accepts both shapes and returns the parsed value (or an
664/// error describing the first mismatch).
665fn parse_llm_json<T: serde::de::DeserializeOwned>(stdout: &str) -> Result<T, String> {
666    // Strategy 1: try the whole stdout as JSON (Claude path).
667    if let Ok(parsed) = serde_json::from_str::<T>(stdout) {
668        return Ok(parsed);
669    }
670    // Strategy 2: walk the JSONL line by line and pick the last
671    // `item.completed` of type `agent_message` (Codex path).
672    let mut last_agent_text: Option<String> = None;
673    for line in stdout.lines() {
674        let line = line.trim();
675        if line.is_empty() {
676            continue;
677        }
678        let Ok(event) = serde_json::from_str::<serde_json::Value>(line) else {
679            continue;
680        };
681        if event.get("type").and_then(|t| t.as_str()) != Some("item.completed") {
682            continue;
683        }
684        let item = match event.get("item") {
685            Some(i) => i,
686            None => continue,
687        };
688        if item.get("type").and_then(|t| t.as_str()) != Some("agent_message") {
689            continue;
690        }
691        if let Some(text) = item.get("text").and_then(|t| t.as_str()) {
692            last_agent_text = Some(text.to_string());
693        }
694    }
695    let text = last_agent_text
696        .ok_or_else(|| "no agent_message found in codex JSONL output".to_string())?;
697    serde_json::from_str::<T>(&text)
698        .map_err(|e| format!("codex agent_message text does not match schema: {e}; raw={text}"))
699}
700
701#[cfg(test)]
702mod tests {
703    use super::*;
704
705    fn test_client(flavour: EmbeddingFlavour, binary: std::path::PathBuf) -> LlmEmbedding {
706        LlmEmbedding {
707            flavour,
708            binary,
709            model: "gpt-5.4".to_string(),
710            codex_schemas: Arc::new(parking_lot::Mutex::new(CodexSchemaFiles::default())),
711        }
712    }
713
714    #[test]
715    #[serial_test::serial(env)]
716    fn oauth_only_enforce_blocks_api_keys() {
717        // SAFETY: this test only sets and unsets env vars; the
718        // `serial(env)` group prevents cross-test interference.
719        unsafe {
720            std::env::set_var("ANTHROPIC_API_KEY", "test");
721            assert!(LlmEmbedding::oauth_only_enforce().is_err());
722            std::env::remove_var("ANTHROPIC_API_KEY");
723
724            std::env::set_var("OPENAI_API_KEY", "test");
725            assert!(LlmEmbedding::oauth_only_enforce().is_err());
726            std::env::remove_var("OPENAI_API_KEY");
727        }
728        assert!(LlmEmbedding::oauth_only_enforce().is_ok());
729    }
730
731    #[test]
732    fn flavour_as_str_is_stable() {
733        assert_eq!(EmbeddingFlavour::Claude.as_str(), "claude");
734        assert_eq!(EmbeddingFlavour::Codex.as_str(), "codex");
735    }
736
737    #[test]
738    fn single_schema_embeds_active_dim() {
739        let schema = build_single_schema(64);
740        assert!(schema.contains(r#""minItems":64"#));
741        assert!(schema.contains(r#""maxItems":64"#));
742        let parsed: serde_json::Value =
743            serde_json::from_str(&schema).expect("single schema must be valid JSON");
744        assert_eq!(parsed["properties"]["embedding"]["minItems"], 64);
745    }
746
747    #[test]
748    fn batch_schema_is_valid_json_and_unbounded_items() {
749        let schema = build_batch_schema(64);
750        let parsed: serde_json::Value =
751            serde_json::from_str(&schema).expect("batch schema must be valid JSON");
752        // The items array must NOT constrain its length so one schema
753        // file serves every batch size (G42/S4).
754        assert!(parsed["properties"]["items"].get("minItems").is_none());
755        assert_eq!(
756            parsed["properties"]["items"]["items"]["properties"]["v"]["minItems"],
757            64
758        );
759    }
760
761    #[test]
762    fn parse_llm_json_accepts_claude_json() {
763        let stdout = r#"{"embedding":[0.0,1.0,2.0]}"#;
764
765        let parsed: EmbeddingResponse = parse_llm_json(stdout).expect("claude JSON must parse");
766
767        assert_eq!(parsed.embedding, vec![0.0, 1.0, 2.0]);
768    }
769
770    #[test]
771    fn parse_llm_json_accepts_codex_jsonl() {
772        let stdout = r#"{"type":"thread.started","thread_id":"mock-thread-0"}
773{"type":"item.completed","item":{"type":"agent_message","text":"{\"embedding\":[0.0,1.0,2.0]}"}}
774{"type":"turn.completed","usage":{"input_tokens":1,"output_tokens":1}}"#;
775
776        let parsed: EmbeddingResponse = parse_llm_json(stdout).expect("codex JSONL must parse");
777
778        assert_eq!(parsed.embedding, vec![0.0, 1.0, 2.0]);
779    }
780
781    #[test]
782    fn parse_llm_json_rejects_jsonl_without_agent_message() {
783        let stdout = r#"{"type":"thread.started","thread_id":"mock-thread-0"}"#;
784
785        let err = parse_llm_json::<EmbeddingResponse>(stdout)
786            .expect_err("missing agent_message must fail");
787
788        assert!(err.contains("no agent_message"));
789    }
790
791    #[test]
792    fn parse_llm_json_accepts_batch_response() {
793        let stdout = r#"{"items":[{"i":1,"v":[0.0,1.0]},{"i":2,"v":[2.0,3.0]}]}"#;
794
795        let parsed: BatchEmbeddingResponse = parse_llm_json(stdout).expect("batch JSON must parse");
796
797        assert_eq!(parsed.items.len(), 2);
798        assert_eq!(parsed.items[0].i, 1);
799        assert_eq!(parsed.items[1].v, vec![2.0, 3.0]);
800    }
801
802    #[test]
803    fn codex_schema_file_is_created_once_and_reused() {
804        let client = test_client(
805            EmbeddingFlavour::Codex,
806            std::path::PathBuf::from("/bin/true"),
807        );
808        let first = client
809            .codex_schema_file(64, false)
810            .expect("schema file must be created");
811        let second = client
812            .codex_schema_file(64, false)
813            .expect("schema file must be reused");
814        assert_eq!(first.path(), second.path(), "same dim must reuse the file");
815
816        let batch = client
817            .codex_schema_file(64, true)
818            .expect("batch schema file must be created");
819        assert_ne!(
820            first.path(),
821            batch.path(),
822            "single and batch schemas are distinct files"
823        );
824
825        let content = std::fs::read_to_string(first.path()).expect("schema file must be readable");
826        assert!(content.contains(r#""minItems":64"#));
827    }
828
829    #[test]
830    fn codex_embedding_command_reads_prompt_from_stdin() {
831        let schema_path = std::env::temp_dir().join("sqlite-graphrag-embed-schema-test.json");
832        let cmd = build_codex_embedding_command(
833            std::path::Path::new("/bin/true"),
834            "gpt-5.4",
835            &schema_path,
836        );
837        let argv: Vec<String> = cmd
838            .as_std()
839            .get_args()
840            .filter_map(|arg| arg.to_str().map(|s| s.to_string()))
841            .collect();
842
843        assert!(
844            argv.iter().any(|arg| arg == "-"),
845            "codex embedding command must read prompt from stdin: {argv:?}"
846        );
847        assert!(
848            !argv.iter().any(|arg| arg.starts_with("passage: ")),
849            "prompt text must not be passed as argv: {argv:?}"
850        );
851        for required in &[
852            "exec",
853            "-c",
854            "sandbox_mode='read-only'",
855            "approval_policy='never'",
856            "--json",
857            "--output-schema",
858            "--ephemeral",
859            "--skip-git-repo-check",
860            "--sandbox",
861            "read-only",
862            "--ignore-user-config",
863            "--ignore-rules",
864            "--model",
865            "gpt-5.4",
866        ] {
867            assert!(
868                argv.iter().any(|arg| arg == required),
869                "missing flag {required} in {argv:?}"
870            );
871        }
872    }
873
874    #[cfg(unix)]
875    #[test]
876    #[serial_test::serial(env)]
877    fn embed_passage_sends_prompt_to_codex_stdin() {
878        use std::os::unix::fs::PermissionsExt;
879
880        // Pin the dimensionality so the mock script and the validation
881        // agree regardless of test execution order.
882        // SAFETY: guarded by serial(env).
883        unsafe {
884            std::env::set_var("SQLITE_GRAPHRAG_EMBEDDING_DIM", "64");
885        }
886
887        let temp = tempfile::tempdir().expect("tempdir must exist");
888        let binary = temp.path().join("codex-stdin-check");
889        let script = r#"#!/usr/bin/env bash
890set -euo pipefail
891
892prompt="$(cat)"
893if [[ "$prompt" != "passage: codex-cli" ]]; then
894  echo "unexpected stdin: $prompt" >&2
895  exit 41
896fi
897
898vals="0.0"
899for _ in $(seq 2 64); do
900  vals="$vals,0.0"
901done
902payload="{\"embedding\":[$vals]}"
903escaped="${payload//\"/\\\"}"
904echo "{\"type\":\"item.completed\",\"item\":{\"type\":\"agent_message\",\"text\":\"$escaped\"}}"
905"#;
906        std::fs::write(&binary, script).expect("mock codex script must be written");
907        let mut perms = std::fs::metadata(&binary)
908            .expect("mock codex metadata must exist")
909            .permissions();
910        perms.set_mode(0o755);
911        std::fs::set_permissions(&binary, perms).expect("mock codex must be executable");
912
913        let embedding = test_client(EmbeddingFlavour::Codex, binary);
914
915        let vector = embedding
916            .embed_passage("codex-cli")
917            .expect("stdin-backed codex embedding must succeed");
918
919        // SAFETY: guarded by serial(env).
920        unsafe {
921            std::env::remove_var("SQLITE_GRAPHRAG_EMBEDDING_DIM");
922        }
923
924        assert_eq!(vector.len(), 64);
925        assert!(vector.iter().all(|value| *value == 0.0));
926    }
927}