Skip to main content

sqz_engine/
codex_integration.rs

1//! OpenAI Codex integration for sqz.
2//!
3//! Codex is OpenAI's terminal coding agent (openai/codex). It does not
4//! expose a stable per-tool-call hook like Claude Code's PreToolUse — the
5//! only event hook documented at the time of writing is `notify`, which
6//! fires at the END of a turn (once per user message) and so is useless
7//! for compressing tool output before the model sees it. An experimental
8//! `features.codex_hooks` flag is mentioned in the config reference but
9//! described as "under development; off by default" and the public `hooks.json`
10//! schema is not documented enough to target safely. See
11//! <https://developers.openai.com/codex/config-reference>.
12//!
13//! What Codex DOES expose reliably today:
14//!
15//!   1. **MCP servers** via `~/.codex/config.toml` under
16//!      `[mcp_servers.<id>]`. The TOML key is `mcp_servers` (snake_case),
17//!      NOT `mcpServers` (camelCase) as in JSON-based tools. See
18//!      <https://github.com/openai/codex/blob/main/docs/config.md>.
19//!
20//!   2. **`AGENTS.md`** — project-level markdown instructions Codex reads
21//!      at session start. It's the Codex analogue of `CLAUDE.md` and
22//!      `.cursor/rules/*.mdc`. A single `AGENTS.md` is the cross-tool
23//!      convention (Codex, GitHub Copilot, Cursor, Windsurf, Amp, Devin
24//!      all read it). See <https://agentsmd.io>.
25//!
26//!   3. **The shell hook** installed by `sqz init` — works transparently
27//!      because Codex runs bash via its sandboxed exec tool and sees the
28//!      compressed stdout automatically. No Codex-specific wiring needed.
29//!
30//! This module implements (1) and (2). The guidance-file approach for (2)
31//! is the standard answer to Codex's "no programmatic hook" limitation.
32
33use std::path::{Path, PathBuf};
34
35use crate::error::{Result, SqzError};
36use crate::kiro_integration::{embedded_sqz_path, SteeringState};
37
38// ── Paths ─────────────────────────────────────────────────────────────────
39
40/// Return the location of Codex's user-level `config.toml`.
41///
42/// Honors `$CODEX_HOME` when set (Codex supports it for isolated installs),
43/// falling back to `~/.codex/config.toml`. Matches the behaviour of the
44/// Codex CLI itself — see the `sqlite_home` discussion in Codex's
45/// `docs/config.md`, which describes the same resolution order.
46pub fn codex_config_path() -> PathBuf {
47    if let Ok(home) = std::env::var("CODEX_HOME") {
48        return PathBuf::from(home).join("config.toml");
49    }
50    let home = std::env::var("HOME")
51        .or_else(|_| std::env::var("USERPROFILE"))
52        .map(PathBuf::from)
53        .unwrap_or_else(|_| PathBuf::from("."));
54    home.join(".codex").join("config.toml")
55}
56
57/// Return the path where Codex reads project-level instructions.
58///
59/// `AGENTS.md` lives in the project root and is the cross-tool standard
60/// (Codex, Copilot, Cursor). We do NOT fall back to the deprecated
61/// `codex.md` / `.codexrules` names — they are pre-consolidation legacy
62/// and new installs should target `AGENTS.md`.
63pub fn agents_md_path(project_dir: &Path) -> PathBuf {
64    project_dir.join("AGENTS.md")
65}
66
67// ── Sentinel-bracketed guidance block ────────────────────────────────────
68//
69// `AGENTS.md` is a free-form markdown file the user may edit freely. sqz
70// writes its guidance between BEGIN/END sentinels so we can later locate
71// and remove exactly what we added, without clobbering unrelated user
72// content. Matches the pattern we already use for shell RC hooks in
73// sqz/src/shell_hook.rs.
74
75const AGENTS_MD_BEGIN: &str = "<!-- BEGIN sqz-guidance (auto-installed by sqz init; remove this block to disable) -->";
76const AGENTS_MD_END: &str = "<!-- END sqz-guidance -->";
77const CODEX_LEGACY_BEGIN: &str = "<!-- BEGIN sqz-agents-guidance (auto-installed by sqz init; remove this block to disable) -->";
78const CODEX_LEGACY_END: &str = "<!-- END sqz-agents-guidance -->";
79const ZED_LEGACY_BEGIN: &str = "<!-- BEGIN sqz-zed-guidance (auto-installed by sqz init; remove this block to disable) -->";
80const ZED_LEGACY_END: &str = "<!-- END sqz-zed-guidance -->";
81
82pub fn sqz_guidance(sqz_path: &str) -> String {
83    format!(
84        "\
85## sqz
86
87Pipe shell commands with long output (git log or diff, tests, builds,
88`ls -la`) through sqz:
89
90```bash
91git status 2>&1 | {sqz_path} compress
92```
93
94Don't pipe interactive commands (REPLs, `ssh`), chains with `&&`, `||`
95or `;`, output redirected to a file, or short output. If that binary is
96missing, run commands without it.
97
98When the sqz MCP tools are available, prefer `sqz_read_file`, `sqz_grep`
99and `sqz_list_dir` over the built-in read, grep and `ls` for long or
100repeated reads, and use `sqz_recall` to find output seen earlier.
101
102A `§ref:HASH§` stands in for output you have already seen, and
103`§ref:HASH:L40-80§` for lines 40-80 of it. Get the text back with the
104`expand` tool or `{sqz_path} expand <token>`. `«A1»` stands for the
105phrase listed for it under `[Abbreviations]` at the top of the output.
106For output without refs or abbreviations, put `SQZ_NO_DEDUP=1` or
107`SQZ_NO_ABBREV=1` right before `{sqz_path} compress`.
108"
109    )
110}
111
112/// Build the markdown block that sqz appends to `AGENTS.md`.
113///
114/// The block is wrapped in HTML comment sentinels so `remove_agents_md_guidance`
115/// can locate and excise it byte-exact on uninstall. `sqz_path` is the
116/// absolute path of the sqz binary; we use it verbatim (no JSON escaping)
117/// because this is a markdown file the user reads — raw backslashes on
118/// Windows render correctly as-is, and the user can copy-paste the
119/// commands straight into their shell.
120pub fn agents_md_guidance_block(sqz_path: &str) -> String {
121    format!("{AGENTS_MD_BEGIN}\n\n{}\n{AGENTS_MD_END}\n", sqz_guidance(sqz_path))
122}
123
124// ── AGENTS.md install/uninstall ──────────────────────────────────────────
125
126#[derive(Debug, Clone, Copy, PartialEq, Eq)]
127pub(crate) enum BlockKind {
128    Shared,
129    Codex,
130    Zed,
131}
132
133impl BlockKind {
134    pub(crate) const ALL: [BlockKind; 3] = [BlockKind::Shared, BlockKind::Codex, BlockKind::Zed];
135
136    fn markers(self) -> (&'static str, &'static str) {
137        match self {
138            BlockKind::Shared => (AGENTS_MD_BEGIN, AGENTS_MD_END),
139            BlockKind::Codex => (CODEX_LEGACY_BEGIN, CODEX_LEGACY_END),
140            BlockKind::Zed => (ZED_LEGACY_BEGIN, ZED_LEGACY_END),
141        }
142    }
143}
144
145#[derive(Debug, Clone, Copy)]
146struct Block {
147    start: usize,
148    end: usize,
149    kind: BlockKind,
150}
151
152fn find_blocks(content: &str) -> Vec<Block> {
153    let mut blocks = Vec::new();
154    for kind in BlockKind::ALL {
155        let (begin, end) = kind.markers();
156        let mut from = 0;
157        while let Some(i) = content[from..].find(begin) {
158            let start = from + i;
159            let Some(j) = content[start..].find(end) else { break };
160            from = start + j + end.len();
161            blocks.push(Block { start, end: from, kind });
162        }
163    }
164    blocks.sort_by_key(|b| b.start);
165    let mut kept: Vec<Block> = Vec::new();
166    for block in blocks {
167        if kept.last().is_some_and(|last| block.start < last.end) {
168            continue;
169        }
170        kept.push(block);
171    }
172    kept
173}
174
175fn has_sqz_marker(content: &str) -> bool {
176    BlockKind::ALL.iter().any(|kind| content.contains(kind.markers().0))
177}
178
179fn excise(content: &str, block: Block) -> String {
180    excise_range(content, block.start, block.end)
181}
182
183pub(crate) fn excise_range(content: &str, start: usize, end: usize) -> String {
184    let head = &content[..start];
185    let after = &content[end..];
186    let tail = after.strip_prefix('\n').unwrap_or(after);
187    let kept_head = head.trim_end_matches('\n');
188    let kept_tail = tail.trim_start_matches('\n');
189    if kept_tail.is_empty() {
190        return if kept_head.is_empty() { String::new() } else { format!("{kept_head}\n") };
191    }
192    if kept_head.is_empty() {
193        return kept_tail.to_string();
194    }
195    let newlines = (head.len() - kept_head.len() + tail.len() - kept_tail.len()).clamp(1, 2);
196    format!("{kept_head}{}{kept_tail}", "\n".repeat(newlines))
197}
198
199fn replace_blocks(content: &str, block: &str) -> Option<String> {
200    let first = *find_blocks(content).first()?;
201    let mut out = format!("{}{block}{}", &content[..first.start], &content[first.end..]);
202    loop {
203        let blocks = find_blocks(&out);
204        let [prev, next, ..] = blocks[..] else { break };
205        out = if out[prev.end..next.start].trim().is_empty() {
206            format!("{}{}", &out[..prev.end], &out[next.end..])
207        } else {
208            excise(&out, next)
209        };
210    }
211    Some(out)
212}
213
214fn remove_blocks(content: &str, drop: &[BlockKind]) -> Option<String> {
215    let mut out = content.to_string();
216    let mut changed = false;
217    while let Some(block) = find_blocks(&out).into_iter().find(|b| drop.contains(&b.kind)) {
218        out = excise(&out, block);
219        changed = true;
220    }
221    changed.then_some(out)
222}
223
224fn append_block(content: &str, block: &str) -> String {
225    if content.trim().is_empty() {
226        return block.to_string();
227    }
228    let mut out = content.to_string();
229    if !out.ends_with('\n') {
230        out.push('\n');
231    }
232    if !out.ends_with("\n\n") {
233        out.push('\n');
234    }
235    out.push_str(block);
236    out
237}
238
239pub(crate) fn replaced_guidance(content: &str, sqz_path: &str) -> Option<String> {
240    replace_blocks(content, agents_md_guidance_block(sqz_path).trim_end_matches('\n'))
241}
242
243pub(crate) fn without_guidance(content: &str) -> Option<String> {
244    remove_blocks(content, &BlockKind::ALL)
245}
246
247fn with_guidance(content: &str, sqz_path: &str) -> Option<String> {
248    let updated = if find_blocks(content).is_empty() {
249        if has_sqz_marker(content) {
250            return None;
251        }
252        append_block(content, &agents_md_guidance_block(sqz_path))
253    } else {
254        replaced_guidance(content, sqz_path)?
255    };
256    (updated != content).then_some(updated)
257}
258
259pub(crate) fn write_guidance(path: &Path, sqz_path: &str) -> Result<bool> {
260    let existing = std::fs::read_to_string(path)
261        .map_err(|e| SqzError::Other(format!("failed to read {}: {e}", path.display())))?;
262    let Some(updated) = with_guidance(&existing, sqz_path) else {
263        return Ok(false);
264    };
265    std::fs::write(path, updated)
266        .map_err(|e| SqzError::Other(format!("failed to write {}: {e}", path.display())))?;
267    Ok(true)
268}
269
270#[derive(Debug, Clone, Copy, PartialEq, Eq)]
271pub enum GuidanceChange {
272    Unchanged,
273    Create,
274    Append,
275    Replace,
276}
277
278pub fn guidance_change(path: &Path, sqz_path: &str) -> GuidanceChange {
279    if !path.exists() {
280        return GuidanceChange::Create;
281    }
282    let Ok(content) = std::fs::read_to_string(path) else {
283        return GuidanceChange::Unchanged;
284    };
285    match with_guidance(&content, sqz_path) {
286        None => GuidanceChange::Unchanged,
287        Some(_) if find_blocks(&content).is_empty() => GuidanceChange::Append,
288        Some(_) => GuidanceChange::Replace,
289    }
290}
291
292#[derive(Debug, Clone, Copy, PartialEq, Eq)]
293pub struct GuidanceBlock {
294    pub bytes: usize,
295    pub legacy: bool,
296}
297
298pub fn guidance_blocks(content: &str) -> Vec<GuidanceBlock> {
299    find_blocks(content)
300        .into_iter()
301        .map(|b| GuidanceBlock { bytes: b.end - b.start, legacy: b.kind != BlockKind::Shared })
302        .collect()
303}
304
305pub fn agents_md_guidance_state(path: &Path) -> SteeringState {
306    match std::fs::read_to_string(path) {
307        Ok(content) => guidance_state(&content),
308        Err(_) => SteeringState::Missing,
309    }
310}
311
312pub(crate) fn guidance_state(content: &str) -> SteeringState {
313    let blocks = find_blocks(content);
314    let [block] = blocks[..] else {
315        return if blocks.is_empty() { SteeringState::UserAuthored } else { SteeringState::Stale };
316    };
317    let text = &content[block.start..block.end];
318    let current = block.kind == BlockKind::Shared
319        && embedded_sqz_path(text).is_some_and(|p| {
320            text.replace("\r\n", "\n") == agents_md_guidance_block(p).trim_end_matches('\n')
321        });
322    if current {
323        SteeringState::Current
324    } else {
325        SteeringState::Stale
326    }
327}
328
329pub fn refresh_agents_md_guidance(path: &Path, fallback_sqz_path: &str) -> Result<bool> {
330    if agents_md_guidance_state(path) != SteeringState::Stale {
331        return Ok(false);
332    }
333    let content = std::fs::read_to_string(path)
334        .map_err(|e| SqzError::Other(format!("failed to read {}: {e}", path.display())))?;
335    let sqz_path = find_blocks(&content)
336        .iter()
337        .filter_map(|b| embedded_sqz_path(&content[b.start..b.end]))
338        .find(|p| Path::new(p).is_file())
339        .unwrap_or(fallback_sqz_path)
340        .to_string();
341    write_guidance(path, &sqz_path)
342}
343
344pub fn codex_mcp_registered() -> bool {
345    std::fs::read_to_string(codex_config_path())
346        .ok()
347        .and_then(|raw| raw.parse::<toml_edit::DocumentMut>().ok())
348        .is_some_and(|doc| {
349            doc.get("mcp_servers")
350                .and_then(|v| v.get("sqz"))
351                .and_then(|v| v.get("command"))
352                .is_some()
353        })
354}
355
356pub(crate) fn remove_guidance_from(
357    path: &Path,
358    drop: &[BlockKind],
359) -> Result<Option<(PathBuf, bool)>> {
360    if !path.exists() {
361        return Ok(None);
362    }
363    let content = std::fs::read_to_string(path)
364        .map_err(|e| SqzError::Other(format!("failed to read {}: {e}", path.display())))?;
365    let Some(updated) = remove_blocks(&content, drop) else {
366        return Ok(Some((path.to_path_buf(), false)));
367    };
368    if is_essentially_empty_agents_md(&updated) {
369        std::fs::remove_file(path)
370            .map_err(|e| SqzError::Other(format!("failed to remove {}: {e}", path.display())))?;
371    } else {
372        std::fs::write(path, updated)
373            .map_err(|e| SqzError::Other(format!("failed to write {}: {e}", path.display())))?;
374    }
375    Ok(Some((path.to_path_buf(), true)))
376}
377
378pub fn remove_all_sqz_guidance(project_dir: &Path) -> Result<Vec<PathBuf>> {
379    let mut changed = Vec::new();
380    for path in [project_dir.join(".rules"), agents_md_path(project_dir)] {
381        if let Some((path, true)) = remove_guidance_from(&path, &BlockKind::ALL)? {
382            changed.push(path);
383        }
384    }
385    Ok(changed)
386}
387
388/// Install sqz's guidance block into `AGENTS.md` at `project_dir`.
389///
390/// If `AGENTS.md` doesn't exist yet, create it with sqz's block as the
391/// sole content. If it exists, append sqz's block (separated by a blank
392/// line so it renders as a new markdown section). If sqz blocks are
393/// already present, the first becomes the current block and the rest
394/// are removed, so `sqz init` stays idempotent.
395///
396/// Returns `true` when the file was created or modified, `false` when
397/// it was left as it was.
398pub fn install_agents_md_guidance(project_dir: &Path, sqz_path: &str) -> Result<bool> {
399    let path = agents_md_path(project_dir);
400    let block = agents_md_guidance_block(sqz_path);
401
402    if path.exists() {
403        return write_guidance(&path, sqz_path);
404    }
405
406    // Fresh AGENTS.md — write just sqz's block. A tiny preamble is added
407    // so `AGENTS.md` is self-explanatory to readers who encounter it for
408    // the first time and don't know the convention.
409    let preamble = "\
410# AGENTS.md
411
412Instructions for AI coding agents (OpenAI Codex, GitHub Copilot, Cursor,
413Windsurf, Amp, Devin) working in this repository. See <https://agentsmd.io>.
414
415";
416    let content = format!("{preamble}{block}");
417    std::fs::write(&path, content)
418        .map_err(|e| SqzError::Other(format!("failed to create {}: {e}", path.display())))?;
419    Ok(true)
420}
421
422/// Remove sqz's guidance block from `AGENTS.md`.
423///
424/// Locates the block by its BEGIN/END sentinels and excises it, keeping
425/// the shared block while Zed is registered and reads this file. If the
426/// resulting `AGENTS.md` is empty (modulo whitespace) or contains only
427/// the stock preamble sqz itself wrote on first install, the file is
428/// deleted entirely — leaving an empty `AGENTS.md` would be noise.
429/// If sentinel markers are missing (user edited them out or the file
430/// was never installed by sqz), this is a no-op.
431///
432/// Returns `Ok((path, removed_or_changed))`:
433///   - `removed_or_changed = true`: the file was modified or deleted
434///   - `removed_or_changed = false`: nothing was removed
435///
436/// Returns `Ok(None)` when `AGENTS.md` does not exist at all.
437pub fn remove_agents_md_guidance(project_dir: &Path) -> Result<Option<(PathBuf, bool)>> {
438    let zed_uses_it = crate::zed_integration::zed_mcp_registered()
439        && crate::zed_integration::zed_guidance_target(project_dir) == agents_md_path(project_dir);
440    remove_agents_md_guidance_at(project_dir, zed_uses_it)
441}
442
443pub(crate) fn remove_agents_md_guidance_at(
444    project_dir: &Path,
445    shared_in_use: bool,
446) -> Result<Option<(PathBuf, bool)>> {
447    let drop: &[BlockKind] = if shared_in_use {
448        &[BlockKind::Codex]
449    } else {
450        &[BlockKind::Codex, BlockKind::Shared]
451    };
452    remove_guidance_from(&agents_md_path(project_dir), drop)
453}
454
455/// Return `true` when the remaining `AGENTS.md` content has no
456/// user-authored material — either empty or just the stock preamble sqz
457/// itself wrote on first install. Used by `remove_agents_md_guidance`
458/// to decide between "write back the trimmed file" and "delete it".
459fn is_essentially_empty_agents_md(s: &str) -> bool {
460    let trimmed = s.trim();
461    if trimmed.is_empty() {
462        return true;
463    }
464    // Match the exact preamble sqz writes on first install (with its
465    // heading, one paragraph of explanation, and no user additions).
466    const PREAMBLE_MARKERS: &[&str] = &[
467        "# AGENTS.md",
468        "Instructions for AI coding agents",
469        "See <https://agentsmd.io>.",
470    ];
471    let has_only_preamble = PREAMBLE_MARKERS.iter().all(|m| trimmed.contains(m))
472        && !trimmed
473            .lines()
474            .any(|l| {
475                let l = l.trim();
476                !l.is_empty()
477                    && !l.starts_with('#')
478                    && !PREAMBLE_MARKERS.iter().any(|m| l.contains(m))
479            });
480    has_only_preamble
481}
482
483// ── ~/.codex/config.toml merger ──────────────────────────────────────────
484
485/// Merge sqz's MCP server entry into Codex's user-level `config.toml`.
486///
487/// Codex reads MCP servers from `~/.codex/config.toml` (or
488/// `$CODEX_HOME/config.toml` when set). The relevant TOML shape is:
489///
490/// ```toml
491/// [mcp_servers.sqz]
492/// command = "sqz-mcp"
493/// args = ["--transport", "stdio"]
494/// ```
495///
496/// Notes:
497///   - Key is `mcp_servers` (snake_case). JSON-tool users who write
498///     `mcpServers` will be quietly ignored by Codex.
499///   - We go through `toml_edit` so the user's existing comments,
500///     key order, and formatting are preserved across the round-trip.
501///     Plain `toml::from_str → toml::to_string` wipes comments.
502///
503/// Idempotent: if `[mcp_servers.sqz]` is already present with a `command`
504/// field, we do not overwrite — the user may have tuned it. Returns
505/// `Ok(false)` in that case.
506///
507/// If the config file doesn't exist yet, it is created with only sqz's
508/// entry (and its parent `~/.codex/` directory is created on demand).
509pub fn install_codex_mcp_config() -> Result<bool> {
510    install_codex_mcp_config_at(None)
511}
512
513/// Internal: home-dir-injectable counterpart used by tests. Avoids
514/// `std::env::set_var` which races with parallel tests that also read
515/// HOME (e.g. the api_proxy property tests that open `~/.sqz/sessions.db`).
516/// See `claude_md_integration::install_claude_mcp_config_at` for the
517/// same pattern.
518pub(crate) fn install_codex_mcp_config_at(home_override: Option<&Path>) -> Result<bool> {
519    let path = match home_override {
520        Some(h) => h.join("config.toml"),
521        None => codex_config_path(),
522    };
523
524    // Read or start from blank. We build on a toml_edit::DocumentMut so
525    // any prior comments/whitespace survive the round-trip.
526    let existing = if path.exists() {
527        std::fs::read_to_string(&path).map_err(|e| {
528            SqzError::Other(format!("failed to read {}: {e}", path.display()))
529        })?
530    } else {
531        String::new()
532    };
533
534    let mut doc: toml_edit::DocumentMut = existing
535        .parse()
536        .map_err(|e| SqzError::Other(format!(
537            "failed to parse {} as TOML: {e}",
538            path.display()
539        )))?;
540
541    // Idempotency: if [mcp_servers.sqz].command is already set, skip.
542    // We only check `command` (not just the table's presence) because a
543    // stub `[mcp_servers.sqz]` with no keys would be a misconfigured
544    // install worth repairing.
545    if let Some(existing_cmd) = doc
546        .get("mcp_servers")
547        .and_then(|v| v.get("sqz"))
548        .and_then(|v| v.get("command"))
549    {
550        if existing_cmd.is_value() {
551            return Ok(false);
552        }
553    }
554
555    // Build the sqz table in place. Using `get_mut(..).or_insert_with(..)`
556    // pattern so we add keys to any existing [mcp_servers] without
557    // replacing the whole table. We deliberately leave the parent
558    // `mcp_servers` implicit (i.e. don't emit an explicit
559    // `[mcp_servers]` header before the sub-tables). The TOML spec
560    // treats `[mcp_servers.sqz]` and `[mcp_servers.jira]` as proper
561    // subtables that implicitly create their parent, and emitting a
562    // bare `[mcp_servers]` header before them is redundant noise in
563    // the file.
564    let mcp_servers = doc
565        .entry("mcp_servers")
566        .or_insert_with(|| {
567            // Start the new table as implicit so toml_edit doesn't
568            // write `[mcp_servers]` before the first subtable header.
569            let mut t = toml_edit::Table::new();
570            t.set_implicit(true);
571            toml_edit::Item::Table(t)
572        })
573        .as_table_mut()
574        .ok_or_else(|| SqzError::Other(format!(
575            "{}: `mcp_servers` is not a table — refusing to overwrite",
576            path.display()
577        )))?;
578
579    let sqz = mcp_servers
580        .entry("sqz")
581        .or_insert_with(|| toml_edit::Item::Table(toml_edit::Table::new()))
582        .as_table_mut()
583        .ok_or_else(|| SqzError::Other(format!(
584            "{}: `mcp_servers.sqz` is not a table — refusing to overwrite",
585            path.display()
586        )))?;
587
588    // Populate the server's required fields. Only write keys the user
589    // hasn't already set so a hand-tuned config survives re-runs.
590    if !sqz.contains_key("command") {
591        sqz["command"] = toml_edit::value("sqz-mcp");
592    }
593    if !sqz.contains_key("args") {
594        let mut args = toml_edit::Array::new();
595        args.push("--transport");
596        args.push("stdio");
597        sqz["args"] = toml_edit::Item::Value(toml_edit::Value::Array(args));
598    }
599
600    // Make sure the parent directory exists.
601    if let Some(parent) = path.parent() {
602        std::fs::create_dir_all(parent).map_err(|e| {
603            SqzError::Other(format!(
604                "failed to create {}: {e}",
605                parent.display()
606            ))
607        })?;
608    }
609
610    std::fs::write(&path, doc.to_string()).map_err(|e| {
611        SqzError::Other(format!("failed to write {}: {e}", path.display()))
612    })?;
613    Ok(true)
614}
615
616/// Remove sqz's MCP entry from Codex's user-level `config.toml`.
617///
618/// Surgical: only removes `[mcp_servers.sqz]`. Other `[mcp_servers.*]`
619/// entries and unrelated keys are preserved. Uses `toml_edit` to keep
620/// the user's comments, key order, and formatting intact.
621///
622/// If removing sqz empties out `[mcp_servers]`, the now-empty table is
623/// dropped so the config file doesn't end up with dangling headers.
624/// If the file would become empty after the surgery (only sqz's entry
625/// existed) we delete it entirely — a zero-byte `config.toml` is just
626/// noise.
627///
628/// Returns `Ok((path, changed))`:
629///   - `changed = true`  — sqz's entry was removed or the file deleted
630///   - `changed = false` — no sqz entry found; nothing changed
631///
632/// Returns `Ok(None)` if the config file does not exist at all.
633pub fn remove_codex_mcp_config() -> Result<Option<(PathBuf, bool)>> {
634    remove_codex_mcp_config_at(None)
635}
636
637/// Internal: home-dir-injectable counterpart used by tests. See
638/// `install_codex_mcp_config_at` for rationale.
639pub(crate) fn remove_codex_mcp_config_at(
640    home_override: Option<&Path>,
641) -> Result<Option<(PathBuf, bool)>> {
642    let path = match home_override {
643        Some(h) => h.join("config.toml"),
644        None => codex_config_path(),
645    };
646    if !path.exists() {
647        return Ok(None);
648    }
649    let raw = std::fs::read_to_string(&path)
650        .map_err(|e| SqzError::Other(format!("failed to read {}: {e}", path.display())))?;
651
652    let mut doc: toml_edit::DocumentMut = match raw.parse() {
653        Ok(d) => d,
654        Err(_) => {
655            // Can't parse — leave it alone rather than destroy user data.
656            return Ok(Some((path, false)));
657        }
658    };
659
660    let mcp_table = match doc.get_mut("mcp_servers").and_then(|v| v.as_table_mut()) {
661        Some(t) => t,
662        None => return Ok(Some((path, false))),
663    };
664    if !mcp_table.contains_key("sqz") {
665        return Ok(Some((path, false)));
666    }
667    mcp_table.remove("sqz");
668
669    // If mcp_servers is now empty, drop the whole table so the file
670    // doesn't carry a dangling `[mcp_servers]` header.
671    let mcp_is_empty = mcp_table.iter().count() == 0;
672    if mcp_is_empty {
673        doc.remove("mcp_servers");
674    }
675
676    // If the whole document is empty now, remove the file.
677    let doc_is_empty = doc.iter().count() == 0;
678    if doc_is_empty {
679        std::fs::remove_file(&path).map_err(|e| {
680            SqzError::Other(format!("failed to remove {}: {e}", path.display()))
681        })?;
682        return Ok(Some((path, true)));
683    }
684
685    std::fs::write(&path, doc.to_string()).map_err(|e| {
686        SqzError::Other(format!("failed to write {}: {e}", path.display()))
687    })?;
688    Ok(Some((path, true)))
689}
690
691// ── Tests ────────────────────────────────────────────────────────────────
692
693#[cfg(test)]
694mod tests {
695    use super::*;
696
697    // codex_config_path() reads CODEX_HOME / HOME env vars. Testing it
698    // with set_var is inherently racy under parallel test execution, so
699    // we only verify the structural invariant: the returned path always
700    // ends with "config.toml".
701    #[test]
702    fn codex_config_path_ends_with_config_toml() {
703        let got = codex_config_path();
704        assert!(
705            got.ends_with("config.toml"),
706            "codex_config_path() must end with config.toml, got: {}",
707            got.display()
708        );
709    }
710
711    // ── AGENTS.md guidance block ─────────────────────────────────────
712
713    #[test]
714    fn agents_md_guidance_block_contains_sqz_invocation() {
715        let block = agents_md_guidance_block("/usr/local/bin/sqz");
716        assert!(block.contains(AGENTS_MD_BEGIN));
717        assert!(block.contains(AGENTS_MD_END));
718        assert!(block.contains("| /usr/local/bin/sqz compress"));
719        assert!(block.contains("sqz_read_file"));
720    }
721
722    #[test]
723    fn shared_guidance_is_short_plain_and_complete() {
724        let text = sqz_guidance("/usr/local/bin/sqz");
725        assert!(text.len() < 1000, "{} bytes", text.len());
726        assert!(!text.contains('%'));
727        assert!(!text.contains('\u{2014}') && !text.contains('\u{2013}'));
728        assert!(text.contains("\ngit status 2>&1 | /usr/local/bin/sqz compress\n"));
729        for needle in [
730            "run commands without it",
731            "sqz_read_file",
732            "sqz_grep",
733            "sqz_list_dir",
734            "sqz_recall",
735            "`expand`",
736            "/usr/local/bin/sqz expand",
737            "§ref:HASH§",
738            "§ref:HASH:L40-80§",
739            "«A1»",
740            "[Abbreviations]",
741            "SQZ_NO_DEDUP=1",
742            "SQZ_NO_ABBREV=1",
743        ] {
744            assert!(text.contains(needle), "missing {needle}");
745        }
746        assert_eq!(embedded_sqz_path(&text), Some("/usr/local/bin/sqz"));
747    }
748
749    fn legacy(kind: &str, body: &str) -> String {
750        format!(
751            "<!-- BEGIN sqz-{kind}-guidance (auto-installed by sqz init; remove this block to disable) -->\n\n{body}\n\n<!-- END sqz-{kind}-guidance -->"
752        )
753    }
754
755    #[test]
756    fn old_blocks_are_replaced_in_place_by_one_block() {
757        let dir = tempfile::tempdir().unwrap();
758        let path = dir.path().join("AGENTS.md");
759        let head = "# Rules\n\nBe polite.\n\n";
760        let tail = "\n\n## Notes\n\nKeep this.\n";
761        std::fs::write(
762            &path,
763            format!(
764                "{head}{}\n\n{}{tail}",
765                legacy("agents", "git status        git status 2>&1 | /old/release/sqz compress"),
766                legacy("zed", "git status 2>&1 | /old/debug/sqz compress"),
767            ),
768        )
769        .unwrap();
770
771        assert!(install_agents_md_guidance(dir.path(), "/new/sqz").unwrap());
772        let block = agents_md_guidance_block("/new/sqz");
773        assert_eq!(
774            std::fs::read_to_string(&path).unwrap(),
775            format!("{head}{}{tail}", block.trim_end_matches('\n'))
776        );
777        assert!(!install_agents_md_guidance(dir.path(), "/new/sqz").unwrap());
778    }
779
780    #[test]
781    fn a_block_between_user_paragraphs_is_removed_without_touching_them() {
782        let dir = tempfile::tempdir().unwrap();
783        let path = dir.path().join("AGENTS.md");
784        std::fs::write(
785            &path,
786            format!(
787                "{}\n\nMine.\n\n{}\n\nAlso mine.\n",
788                legacy("zed", "git status 2>&1 | /old/sqz compress"),
789                legacy("agents", "git status        git status 2>&1 | /old/sqz compress"),
790            ),
791        )
792        .unwrap();
793
794        assert!(install_agents_md_guidance(dir.path(), "/new/sqz").unwrap());
795        let block = agents_md_guidance_block("/new/sqz");
796        assert_eq!(
797            std::fs::read_to_string(&path).unwrap(),
798            format!("{}\n\nMine.\n\nAlso mine.\n", block.trim_end_matches('\n'))
799        );
800    }
801
802    #[test]
803    fn refresh_keeps_the_first_binary_that_still_exists() {
804        let dir = tempfile::tempdir().unwrap();
805        let bin = dir.path().join("sqz");
806        std::fs::write(&bin, "").unwrap();
807        let bin = bin.to_string_lossy().to_string();
808        let path = dir.path().join("AGENTS.md");
809        std::fs::write(
810            &path,
811            format!(
812                "{}\n\n{}\n",
813                legacy("agents", "git status        git status 2>&1 | /gone/sqz compress"),
814                legacy("zed", &format!("git status 2>&1 | {bin} compress")),
815            ),
816        )
817        .unwrap();
818
819        assert_eq!(agents_md_guidance_state(&path), SteeringState::Stale);
820        assert!(refresh_agents_md_guidance(&path, "/fallback/sqz").unwrap());
821        assert_eq!(std::fs::read_to_string(&path).unwrap(), agents_md_guidance_block(&bin));
822        assert_eq!(agents_md_guidance_state(&path), SteeringState::Current);
823        assert!(!refresh_agents_md_guidance(&path, "/fallback/sqz").unwrap());
824    }
825
826    #[test]
827    fn refresh_falls_back_when_the_old_binary_is_gone() {
828        let dir = tempfile::tempdir().unwrap();
829        let path = dir.path().join("AGENTS.md");
830        std::fs::write(
831            &path,
832            format!("{}\n", legacy("agents", "git status        git status 2>&1 | /gone/sqz compress")),
833        )
834        .unwrap();
835
836        assert!(refresh_agents_md_guidance(&path, "/fallback/sqz").unwrap());
837        assert_eq!(std::fs::read_to_string(&path).unwrap(), agents_md_guidance_block("/fallback/sqz"));
838    }
839
840    #[test]
841    fn outdated_text_in_the_new_block_is_stale() {
842        let dir = tempfile::tempdir().unwrap();
843        let path = dir.path().join("AGENTS.md");
844        std::fs::write(&path, agents_md_guidance_block("/bin/sqz").replace("sqz_recall", "recall")).unwrap();
845        assert_eq!(agents_md_guidance_state(&path), SteeringState::Stale);
846        std::fs::write(&path, format!("{0}\n{0}", agents_md_guidance_block("/bin/sqz"))).unwrap();
847        assert_eq!(agents_md_guidance_state(&path), SteeringState::Stale);
848    }
849
850    #[test]
851    fn files_without_a_complete_sqz_block_are_left_alone() {
852        let dir = tempfile::tempdir().unwrap();
853        let path = dir.path().join("AGENTS.md");
854        std::fs::write(&path, "# Mine\n").unwrap();
855        assert_eq!(agents_md_guidance_state(&path), SteeringState::UserAuthored);
856        assert!(!refresh_agents_md_guidance(&path, "/fallback/sqz").unwrap());
857        assert_eq!(std::fs::read_to_string(&path).unwrap(), "# Mine\n");
858
859        let broken = "# Mine\n\n<!-- BEGIN sqz-agents-guidance (auto-installed by sqz init; remove this block to disable) -->\nhalf\n";
860        std::fs::write(&path, broken).unwrap();
861        assert!(!install_agents_md_guidance(dir.path(), "/new/sqz").unwrap());
862        assert_eq!(std::fs::read_to_string(&path).unwrap(), broken);
863    }
864
865    #[test]
866    fn guidance_change_says_what_init_would_do() {
867        let dir = tempfile::tempdir().unwrap();
868        let path = dir.path().join("AGENTS.md");
869        assert_eq!(guidance_change(&path, "/bin/sqz"), GuidanceChange::Create);
870        std::fs::write(&path, "# Mine\n").unwrap();
871        assert_eq!(guidance_change(&path, "/bin/sqz"), GuidanceChange::Append);
872        std::fs::write(&path, format!("# Mine\n\n{}\n", legacy("zed", "git status 2>&1 | /bin/sqz compress"))).unwrap();
873        assert_eq!(guidance_change(&path, "/bin/sqz"), GuidanceChange::Replace);
874        install_agents_md_guidance(dir.path(), "/bin/sqz").unwrap();
875        assert_eq!(guidance_change(&path, "/bin/sqz"), GuidanceChange::Unchanged);
876    }
877
878    #[test]
879    fn codex_removal_keeps_the_shared_block_while_zed_uses_it() {
880        let dir = tempfile::tempdir().unwrap();
881        let path = dir.path().join("AGENTS.md");
882        let block = agents_md_guidance_block("sqz");
883        std::fs::write(&path, format!("# Mine\n\n{}\n\n{block}", legacy("agents", "a"))).unwrap();
884
885        let (_, changed) = remove_agents_md_guidance_at(dir.path(), true).unwrap().unwrap();
886        assert!(changed);
887        assert_eq!(std::fs::read_to_string(&path).unwrap(), format!("# Mine\n\n{block}"));
888        let (_, changed) = remove_agents_md_guidance_at(dir.path(), true).unwrap().unwrap();
889        assert!(!changed);
890        let (_, changed) = remove_agents_md_guidance_at(dir.path(), false).unwrap().unwrap();
891        assert!(changed);
892        assert_eq!(std::fs::read_to_string(&path).unwrap(), "# Mine\n");
893    }
894
895    #[test]
896    fn full_uninstall_removes_old_and_new_blocks() {
897        let dir = tempfile::tempdir().unwrap();
898        let agents = dir.path().join("AGENTS.md");
899        let rules = dir.path().join(".rules");
900        std::fs::write(
901            &agents,
902            format!("# Mine\n\n{}\n\n{}", legacy("agents", "a"), agents_md_guidance_block("sqz")),
903        )
904        .unwrap();
905        std::fs::write(&rules, format!("Tabs.\n\n{}\n", legacy("zed", "z"))).unwrap();
906
907        let changed = remove_all_sqz_guidance(dir.path()).unwrap();
908        assert_eq!(changed, vec![rules.clone(), agents.clone()]);
909        assert_eq!(std::fs::read_to_string(&agents).unwrap(), "# Mine\n");
910        assert_eq!(std::fs::read_to_string(&rules).unwrap(), "Tabs.\n");
911    }
912
913    #[test]
914    fn install_agents_md_creates_file_with_preamble() {
915        let dir = tempfile::tempdir().unwrap();
916        let created = install_agents_md_guidance(dir.path(), "sqz").unwrap();
917        assert!(created);
918        let content = std::fs::read_to_string(dir.path().join("AGENTS.md")).unwrap();
919        assert!(content.starts_with("# AGENTS.md"));
920        assert!(content.contains(AGENTS_MD_BEGIN));
921        assert!(content.contains(AGENTS_MD_END));
922    }
923
924    #[test]
925    fn install_agents_md_appends_without_clobbering_user_content() {
926        let dir = tempfile::tempdir().unwrap();
927        let path = dir.path().join("AGENTS.md");
928        std::fs::write(
929            &path,
930            "# My project rules\n\nBe polite. Run tests before committing.\n",
931        ).unwrap();
932
933        let created = install_agents_md_guidance(dir.path(), "sqz").unwrap();
934        assert!(created);
935
936        let content = std::fs::read_to_string(&path).unwrap();
937        assert!(content.contains("# My project rules"),
938            "original heading must survive");
939        assert!(content.contains("Be polite. Run tests before committing."),
940            "original body must survive");
941        let user_idx = content.find("Be polite").unwrap();
942        let sqz_idx = content.find(AGENTS_MD_BEGIN).unwrap();
943        assert!(sqz_idx > user_idx,
944            "sqz's block must append after user content, not prepend");
945    }
946
947    #[test]
948    fn install_agents_md_is_idempotent() {
949        let dir = tempfile::tempdir().unwrap();
950        let first = install_agents_md_guidance(dir.path(), "sqz").unwrap();
951        assert!(first);
952        let second = install_agents_md_guidance(dir.path(), "sqz").unwrap();
953        assert!(!second, "second install must be a no-op");
954
955        let content = std::fs::read_to_string(dir.path().join("AGENTS.md")).unwrap();
956        let occurrences = content.matches(AGENTS_MD_BEGIN).count();
957        assert_eq!(occurrences, 1, "must not duplicate the block on re-install");
958    }
959
960    #[test]
961    fn remove_agents_md_preserves_user_content() {
962        let dir = tempfile::tempdir().unwrap();
963        let path = dir.path().join("AGENTS.md");
964        std::fs::write(
965            &path,
966            "# My project rules\n\nBe polite. Run tests before committing.\n",
967        ).unwrap();
968        install_agents_md_guidance(dir.path(), "sqz").unwrap();
969
970        let (returned_path, changed) =
971            remove_agents_md_guidance_at(dir.path(), false).unwrap().unwrap();
972        assert_eq!(returned_path, path);
973        assert!(changed);
974        assert!(path.exists(),
975            "file must NOT be deleted — it has user content");
976
977        let content = std::fs::read_to_string(&path).unwrap();
978        assert!(!content.contains(AGENTS_MD_BEGIN));
979        assert!(!content.contains(AGENTS_MD_END));
980        assert!(content.contains("# My project rules"),
981            "user heading must survive the uninstall");
982        assert!(content.contains("Be polite. Run tests before committing."),
983            "user body must survive the uninstall");
984    }
985
986    #[test]
987    fn remove_agents_md_deletes_file_when_only_sqz_preamble_remains() {
988        let dir = tempfile::tempdir().unwrap();
989        install_agents_md_guidance(dir.path(), "sqz").unwrap();
990        let path = dir.path().join("AGENTS.md");
991        assert!(path.exists());
992
993        let (_returned, changed) =
994            remove_agents_md_guidance_at(dir.path(), false).unwrap().unwrap();
995        assert!(changed);
996        assert!(
997            !path.exists(),
998            "fresh-install AGENTS.md must be removed when sqz block is stripped"
999        );
1000    }
1001
1002    #[test]
1003    fn remove_agents_md_noop_when_block_missing() {
1004        let dir = tempfile::tempdir().unwrap();
1005        let path = dir.path().join("AGENTS.md");
1006        std::fs::write(&path, "# User-authored, sqz never touched this.\n").unwrap();
1007
1008        let (returned_path, changed) =
1009            remove_agents_md_guidance_at(dir.path(), false).unwrap().unwrap();
1010        assert_eq!(returned_path, path);
1011        assert!(!changed, "no sqz block means no change");
1012        assert!(path.exists(), "user file must be untouched");
1013    }
1014
1015    #[test]
1016    fn remove_agents_md_returns_none_when_file_missing() {
1017        let dir = tempfile::tempdir().unwrap();
1018        let result = remove_agents_md_guidance_at(dir.path(), false).unwrap();
1019        assert!(result.is_none());
1020    }
1021
1022    // ── ~/.codex/config.toml merger ──────────────────────────────────
1023
1024    #[test]
1025    fn install_codex_mcp_config_creates_file_with_sqz_entry() {
1026        let dir = tempfile::tempdir().unwrap();
1027        let created = install_codex_mcp_config_at(Some(dir.path())).unwrap();
1028        assert!(created);
1029        let content = std::fs::read_to_string(dir.path().join("config.toml")).unwrap();
1030        assert!(
1031            content.contains("[mcp_servers.sqz]"),
1032            "config.toml must contain [mcp_servers.sqz] header; got:\n{content}"
1033        );
1034        assert!(content.contains("command = \"sqz-mcp\""),
1035            "command must be sqz-mcp");
1036        assert!(content.contains("--transport"));
1037        assert!(content.contains("stdio"));
1038    }
1039
1040    #[test]
1041    fn install_codex_mcp_config_preserves_existing_other_servers() {
1042        let dir = tempfile::tempdir().unwrap();
1043        let cfg = dir.path().join("config.toml");
1044        std::fs::write(
1045            &cfg,
1046            "# User's existing Codex config, with a comment.\n\
1047             model = \"gpt-5\"\n\
1048             \n\
1049             [mcp_servers.other]\n\
1050             command = \"other-server\"\n\
1051             args = [\"--flag\"]\n",
1052        ).unwrap();
1053
1054        let created = install_codex_mcp_config_at(Some(dir.path())).unwrap();
1055        assert!(created);
1056
1057        let after = std::fs::read_to_string(&cfg).unwrap();
1058        assert!(after.contains("# User's existing Codex config"),
1059            "comment must survive: {after}");
1060        assert!(after.contains("model = \"gpt-5\""),
1061            "top-level key must survive: {after}");
1062        assert!(after.contains("[mcp_servers.other]"),
1063            "existing server entry must survive: {after}");
1064        assert!(after.contains("command = \"other-server\""),
1065            "existing server command must survive: {after}");
1066        assert!(after.contains("[mcp_servers.sqz]"),
1067            "sqz entry must be added: {after}");
1068    }
1069
1070    #[test]
1071    fn install_codex_mcp_config_is_idempotent() {
1072        let dir = tempfile::tempdir().unwrap();
1073        assert!(install_codex_mcp_config_at(Some(dir.path())).unwrap());
1074        assert!(
1075            !install_codex_mcp_config_at(Some(dir.path())).unwrap(),
1076            "second install with complete [mcp_servers.sqz] must be a no-op"
1077        );
1078    }
1079
1080    #[test]
1081    fn install_codex_mcp_config_does_not_overwrite_user_tuned_entry() {
1082        let dir = tempfile::tempdir().unwrap();
1083        let cfg = dir.path().join("config.toml");
1084        std::fs::write(
1085            &cfg,
1086            "[mcp_servers.sqz]\n\
1087             command = \"/custom/path/sqz-mcp\"\n\
1088             args = [\"--transport\", \"sse\", \"--port\", \"3999\"]\n",
1089        ).unwrap();
1090
1091        let changed = install_codex_mcp_config_at(Some(dir.path())).unwrap();
1092        assert!(!changed, "existing complete entry must be idempotent-skipped");
1093        let after = std::fs::read_to_string(&cfg).unwrap();
1094        assert!(after.contains("/custom/path/sqz-mcp"),
1095            "user's custom command must survive re-init");
1096        assert!(after.contains("\"sse\""),
1097            "user's custom transport must survive");
1098    }
1099
1100    #[test]
1101    fn remove_codex_mcp_config_removes_only_sqz_entry() {
1102        let dir = tempfile::tempdir().unwrap();
1103        let cfg = dir.path().join("config.toml");
1104        std::fs::write(
1105            &cfg,
1106            "# keep this comment\n\
1107             model = \"gpt-5\"\n\
1108             \n\
1109             [mcp_servers.other]\n\
1110             command = \"other-server\"\n\
1111             \n\
1112             [mcp_servers.sqz]\n\
1113             command = \"sqz-mcp\"\n\
1114             args = [\"--transport\", \"stdio\"]\n",
1115        ).unwrap();
1116
1117        let (path, changed) = remove_codex_mcp_config_at(Some(dir.path())).unwrap().unwrap();
1118        assert_eq!(path, cfg);
1119        assert!(changed);
1120
1121        let after = std::fs::read_to_string(&cfg).unwrap();
1122        assert!(after.contains("# keep this comment"),
1123            "comment must survive: {after}");
1124        assert!(after.contains("model = \"gpt-5\""),
1125            "top-level key must survive: {after}");
1126        assert!(after.contains("[mcp_servers.other]"),
1127            "other server entry must survive: {after}");
1128        assert!(!after.contains("[mcp_servers.sqz]"),
1129            "sqz entry must be gone: {after}");
1130    }
1131
1132    #[test]
1133    fn remove_codex_mcp_config_deletes_file_when_sqz_was_the_only_entry() {
1134        let dir = tempfile::tempdir().unwrap();
1135        let cfg = dir.path().join("config.toml");
1136        install_codex_mcp_config_at(Some(dir.path())).unwrap();
1137        let (path, changed) = remove_codex_mcp_config_at(Some(dir.path())).unwrap().unwrap();
1138        assert_eq!(path, cfg);
1139        assert!(changed);
1140        assert!(!cfg.exists(),
1141            "config.toml with only sqz must be deleted on uninstall");
1142    }
1143
1144    #[test]
1145    fn remove_codex_mcp_config_returns_none_when_file_missing() {
1146        let dir = tempfile::tempdir().unwrap();
1147        let result = remove_codex_mcp_config_at(Some(dir.path())).unwrap();
1148        assert!(result.is_none());
1149    }
1150
1151    #[test]
1152    fn remove_codex_mcp_config_noop_when_sqz_entry_missing() {
1153        let dir = tempfile::tempdir().unwrap();
1154        let cfg = dir.path().join("config.toml");
1155        std::fs::write(&cfg, "[mcp_servers.other]\ncommand = \"x\"\n").unwrap();
1156        let (path, changed) = remove_codex_mcp_config_at(Some(dir.path())).unwrap().unwrap();
1157        assert_eq!(path, cfg);
1158        assert!(!changed);
1159        let after = std::fs::read_to_string(&cfg).unwrap();
1160        assert!(after.contains("[mcp_servers.other]"));
1161    }
1162}