Skip to main content

ai_memory/cli/
wrap.rs

1// Copyright 2026 AlphaOne LLC
2// SPDX-License-Identifier: Apache-2.0
3
4//! `ai-memory wrap <agent>` — cross-platform Rust replacement for the
5//! shell wrappers PR-1 of issue #487 shipped in the integration recipes.
6//!
7//! ## What it does
8//!
9//! 1. Calls `cli::boot::run` in-process, capturing its stdout into a
10//!    buffer. No subprocess; no shell. The `--no-boot` flag skips this
11//!    step so a misconfigured DB path doesn't block the agent.
12//! 2. Builds a system-context string of the form
13//!    `<preamble>\n\n<boot output>` where the preamble explains to the
14//!    downstream agent that it has ai-memory access.
15//! 3. Spawns the wrapped agent (`std::process::Command`) with the
16//!    system-context delivered via the chosen strategy:
17//!    - `SystemFlag` — `<agent> <flag> "<system_msg>" <trailing args...>`
18//!    - `SystemEnv`  — `<env_name>=<system_msg> <agent> <trailing args...>`
19//!    - `MessageFile` — write `<system_msg>` to a `NamedTempFile`, pass
20//!      `<flag> <tempfile_path>` to the agent, drop the tempfile on
21//!      exit so it is cleaned up by the OS.
22//!    - `Auto` — resolved at runtime from a built-in lookup table
23//!      (`default_strategy`).
24//! 4. Forwards the parent's stdin / stdout / stderr unmodified
25//!    (`Stdio::inherit`).
26//! 5. Returns the wrapped agent's exit code as the wrap subcommand's
27//!    exit code, so wrappers compose cleanly with shell pipelines and
28//!    CI gates that branch on `$?`.
29//!
30//! ## Why Rust, not bash + PowerShell
31//!
32//! The user directive on issue #487 PR-6 was: implementation should be
33//! predominantly Rust with config hooks. PR-1 shipped per-recipe bash
34//! and PowerShell wrappers, which doubled the maintenance surface and
35//! couldn't run in restricted Windows / containerized environments
36//! without a shell. A single cross-platform Rust subcommand eliminates
37//! both problems — it's the same code path on macOS / Linux / Windows
38//! / Docker / Kubernetes / Nix / etc.
39//!
40//! ## Lookup table
41//!
42//! [`crate::llm_cli_wrap::default_strategy`] resolves the unflagged
43//! form `ai-memory wrap <agent> -- <args>` to the right delivery
44//! mechanism for the agents we can identify by name today. Unknown
45//! agents fall through to `--system <msg>` because that's the most
46//! common contract across OpenAI-compatible CLIs. Future PRs (notably
47//! PR-7) can extend the table by adding match arms.
48//!
49//! ## Substrate split (#1183)
50//!
51//! The per-CLI-binary `WrapStrategy` enum + the per-vendor table live
52//! in [`crate::llm_cli_wrap`], adjacent to [`crate::llm`]'s alias
53//! tables, so the per-vendor substrate has one home per concern (HTTP
54//! wire shape in `llm.rs`, CLI ABI in `llm_cli_wrap.rs`). The
55//! CLI-binary-name detection logic that PICKS a `WrapStrategy` stays
56//! HERE because it's CLI-specific (clap `WrapArgs` overrides → table
57//! fallback).
58
59use crate::cli::CliOutput;
60use crate::cli::boot::{self, BootArgs};
61use crate::llm_cli_wrap::{WrapStrategy, default_strategy};
62use anyhow::{Context, Result};
63use clap::Args;
64use std::ffi::OsStr;
65use std::io::Write;
66use std::path::Path;
67use std::process::{Command, Stdio};
68
69/// Default budget for the inner `ai-memory boot` call when the caller
70/// doesn't override. Mirrors `cli::boot::DEFAULT_BUDGET_TOKENS` but is
71/// re-declared here so wrap can tune independently if needed.
72const DEFAULT_WRAP_BUDGET_TOKENS: usize = 4096;
73
74/// Default row limit for the inner boot call. Same value `cli::boot`
75/// itself defaults to.
76const DEFAULT_WRAP_LIMIT: usize = 10;
77
78/// Preamble injected before the boot output in every wrap call.
79/// Explains to the downstream agent why it's seeing this context. Kept
80/// short and stable so prompt-cache breakpoints upstream stay warm.
81const WRAP_PREAMBLE: &str = "You have access to ai-memory, a persistent memory system. \
82The recent context loaded for you appears below. Reference it when relevant to the user's request.";
83
84/// #1575 — subdirectory of the per-user ai-memory data dir
85/// (`~/.ai-memory`, [`crate::AI_MEMORY_HOME_DIR_NAME`]) where the
86/// `MessageFile` strategy stages the boot-context system message.
87const WRAP_STAGING_SUBDIR: &str = "wrap";
88
89/// #1575 — resolve (and secure) the staging directory for the
90/// `MessageFile` boot-context file: `~/.ai-memory/wrap/`, mode 0700.
91///
92/// The boot-context system message contains memory contents, so it
93/// must not sit on a world-readable tmpfs path for the wrapped
94/// agent's whole lifetime (the pre-#1575 behavior — `NamedTempFile`
95/// under `std::env::temp_dir()`). Returns `None` when the home
96/// directory cannot be resolved or the directory cannot be created /
97/// permission-tightened; the caller then falls back to the platform
98/// temp dir with an operator-visible WARN.
99fn message_file_staging_dir() -> Option<std::path::PathBuf> {
100    let dir = dirs::home_dir()?
101        .join(crate::AI_MEMORY_HOME_DIR_NAME)
102        .join(WRAP_STAGING_SUBDIR);
103    std::fs::create_dir_all(&dir).ok()?;
104    #[cfg(unix)]
105    {
106        use std::os::unix::fs::PermissionsExt;
107        std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o700)).ok()?;
108    }
109    Some(dir)
110}
111
112/// Args for `ai-memory wrap`. Designed so the simplest form
113/// (`ai-memory wrap codex -- "hello"`) just works — every flag has a
114/// defaulted value or the lookup table fills it in.
115#[derive(Args, Debug)]
116pub struct WrapArgs {
117    /// Name of the agent CLI to wrap, e.g. `codex`, `aider`, `gemini`,
118    /// `ollama`. Resolved against
119    /// [`crate::llm_cli_wrap::default_strategy`] to pick the
120    /// system-message delivery mechanism unless the user overrides
121    /// with one of the strategy flags below. The agent name is also
122    /// the executable looked up on `$PATH`.
123    pub agent: String,
124
125    /// Override the system-message flag (e.g. `--system-prompt`). When
126    /// set, wrap delivers the system message via this flag regardless
127    /// of what the lookup table says for `<agent>`.
128    #[arg(long, value_name = "FLAG")]
129    pub system_flag: Option<String>,
130
131    /// Override the system-message env var (e.g. `OPENAI_CLI_SYSTEM`).
132    /// Mutually exclusive with `--system-flag` and
133    /// `--message-file-flag`; if multiple are set, the last specified
134    /// on the command line wins (clap default), but the most common
135    /// case is supplying exactly one.
136    #[arg(long, value_name = "NAME", conflicts_with_all = ["system_flag", "message_file_flag"])]
137    pub system_env: Option<String>,
138
139    /// Override the message-file flag (e.g. `--message-file`). Wrap
140    /// will write the system message to a tempfile and pass this flag
141    /// + the tempfile path to the agent. The tempfile is cleaned up on
142    /// wrap exit (cross-platform; uses `tempfile::NamedTempFile`).
143    #[arg(long, value_name = "FLAG", conflicts_with_all = ["system_flag", "system_env"])]
144    pub message_file_flag: Option<String>,
145
146    /// Skip the inner `ai-memory boot` call entirely. The wrapped
147    /// agent runs without any prepended memory context. Useful when
148    /// the DB is known to be unavailable, when the user wants the wrap
149    /// subcommand for argv-forwarding only, or for tests that want to
150    /// isolate the wrapping behavior from the boot-loading behavior.
151    #[arg(long, default_value_t = false)]
152    pub no_boot: bool,
153
154    /// Row limit forwarded to the inner `ai-memory boot --limit`.
155    /// Clamped to `[1, 50]` by `cli::boot` itself.
156    #[arg(long, default_value_t = DEFAULT_WRAP_LIMIT)]
157    pub limit: usize,
158
159    /// Approximate token budget forwarded to the inner
160    /// `ai-memory boot --budget-tokens`.
161    #[arg(long, default_value_t = DEFAULT_WRAP_BUDGET_TOKENS)]
162    pub budget_tokens: usize,
163
164    /// Trailing arguments forwarded verbatim to the wrapped agent CLI
165    /// after the system-message delivery (the convention is to
166    /// separate them with `--` on the command line:
167    /// `ai-memory wrap codex -- chat --model gpt-5`).
168    #[arg(trailing_var_arg = true, allow_hyphen_values = true)]
169    pub trailing: Vec<String>,
170}
171
172/// Resolve the active strategy from the user-supplied overrides plus
173/// the built-in lookup table. Order of precedence:
174///
175/// 1. `--system-env <name>` → `SystemEnv`
176/// 2. `--message-file-flag <flag>` → `MessageFile`
177/// 3. `--system-flag <flag>` → `SystemFlag`
178/// 4. fall through to
179///    [`crate::llm_cli_wrap::default_strategy`]`(agent)` (the
180///    per-CLI-binary lookup table)
181fn resolve_strategy(args: &WrapArgs) -> WrapStrategy {
182    if let Some(name) = args.system_env.as_deref() {
183        return WrapStrategy::SystemEnv { name: name.into() };
184    }
185    if let Some(flag) = args.message_file_flag.as_deref() {
186        return WrapStrategy::MessageFile { flag: flag.into() };
187    }
188    if let Some(flag) = args.system_flag.as_deref() {
189        return WrapStrategy::SystemFlag { flag: flag.into() };
190    }
191    default_strategy(&args.agent)
192}
193
194/// Run `cli::boot::run` in-process, capturing its stdout into a
195/// `Vec<u8>`. Stderr is also captured but discarded — the boot helper
196/// already honors `--quiet` for us, so any stderr that escapes is by
197/// design (a developer-facing diagnostic).
198///
199/// On any boot failure, this function returns an empty `String` rather
200/// than propagating — the agent should still run even if memory load
201/// fails. The user-facing diagnostic header is already on stdout in
202/// that case (`# ai-memory boot: warn — db unavailable …`) so the
203/// caller still sees what happened.
204fn run_boot_capture(
205    db_path: &Path,
206    limit: usize,
207    budget_tokens: usize,
208    app_config: &crate::config::AppConfig,
209) -> String {
210    let mut stdout: Vec<u8> = Vec::new();
211    let mut stderr: Vec<u8> = Vec::new();
212    let mut out = CliOutput::from_std(&mut stdout, &mut stderr);
213    let args = BootArgs {
214        namespace: None,
215        limit,
216        budget_tokens,
217        format: "text".to_string(),
218        no_header: false,
219        // --quiet so a missing DB never blocks the wrapped agent.
220        quiet: true,
221        cwd: None,
222    };
223    if boot::run(db_path, &args, app_config, &mut out).is_err() {
224        // Even on hard failure (which `cli::boot::run` should never
225        // hit thanks to the `--quiet` graceful path), return an empty
226        // string so the agent runs unwrapped rather than getting a
227        // blocking error.
228        return String::new();
229    }
230    String::from_utf8(stdout).unwrap_or_default()
231}
232
233/// Assemble the `<preamble>\n\n<boot_output>` system message. Trims
234/// trailing whitespace on the boot section to keep the assembled
235/// string tidy in the agent's prompt.
236fn build_system_message(boot_output: &str) -> String {
237    let trimmed = boot_output.trim_end();
238    if trimmed.is_empty() {
239        // Even with an empty body the preamble is still useful — it
240        // tells the agent "you have memory access" so it knows it can
241        // call `memory_recall` mid-session if it has the tool.
242        WRAP_PREAMBLE.to_string()
243    } else {
244        format!("{WRAP_PREAMBLE}\n\n{trimmed}")
245    }
246}
247
248/// Spawn the agent with stdio inherited and return the exit code.
249/// Wrapped here so tests can assert on the spawned-command shape via
250/// the helpers in `#[cfg(test)] mod tests`.
251fn spawn_and_wait(mut cmd: Command) -> Result<i32> {
252    cmd.stdin(Stdio::inherit())
253        .stdout(Stdio::inherit())
254        .stderr(Stdio::inherit());
255    let status = cmd
256        .status()
257        .with_context(|| format!("ai-memory wrap: failed to spawn agent {cmd:?}"))?;
258    // Unix: `code()` is None when the child was killed by a signal.
259    // We then surface 128+sig per the standard shell convention so the
260    // caller can branch on the signal in CI scripts.
261    let code = if let Some(c) = status.code() {
262        c
263    } else {
264        #[cfg(unix)]
265        {
266            use std::os::unix::process::ExitStatusExt;
267            status.signal().map_or(1, |s| 128 + s)
268        }
269        #[cfg(not(unix))]
270        {
271            1
272        }
273    };
274    Ok(code)
275}
276
277/// Build the `Command` for an agent given a strategy. Pulled out of
278/// `run` so the tests can assert directly on the resulting `Command`'s
279/// argv / env without spawning a subprocess.
280///
281/// Returns the assembled `Command` + (when the strategy is
282/// `MessageFile`) the `NamedTempFile` whose lifetime governs cleanup.
283/// The caller MUST keep the returned `Option<NamedTempFile>` alive
284/// until after the child has exited; dropping it sooner unlinks the
285/// file mid-spawn on platforms where unlink-while-open is permitted.
286fn build_command_for_strategy(
287    agent: &str,
288    strategy: &WrapStrategy,
289    system_msg: &str,
290    trailing: &[String],
291) -> Result<(Command, Option<tempfile::NamedTempFile>)> {
292    let mut cmd = Command::new(agent);
293    let mut tempfile_handle: Option<tempfile::NamedTempFile> = None;
294    match strategy {
295        WrapStrategy::SystemFlag { flag } => {
296            cmd.arg(flag).arg(system_msg);
297            for t in trailing {
298                cmd.arg(t);
299            }
300        }
301        WrapStrategy::SystemEnv { name } => {
302            cmd.env(name, system_msg);
303            for t in trailing {
304                cmd.arg(t);
305            }
306        }
307        WrapStrategy::MessageFile { flag } => {
308            // `tempfile::NamedTempFile` is cross-platform: on Unix it's
309            // a regular file with a randomised name; on Windows it
310            // skips the unlink-while-open trick (which Windows
311            // disallows) and cleans up on `Drop`. Either way the file
312            // is gone after wrap exits.
313            //
314            // #1575 — stage under `~/.ai-memory/wrap/` (0700 dir,
315            // 0600 file) instead of the platform temp dir, so the
316            // memory-bearing boot context never sits on a
317            // world-readable tmpfs path for the agent's lifetime.
318            // The temp dir remains ONLY as a home-unresolvable
319            // fallback, with an operator-visible WARN.
320            let mut tf = match message_file_staging_dir() {
321                Some(dir) => tempfile::NamedTempFile::new_in(&dir).context(
322                    "ai-memory wrap: failed to create system-message file in staging dir",
323                )?,
324                None => {
325                    tracing::warn!(
326                        "ai-memory wrap: could not resolve/secure the {}/{} staging dir under \
327                         the home directory; falling back to the platform temp dir for the \
328                         boot-context message file (#1575 — memory contents will transit a \
329                         shared temp path)",
330                        crate::AI_MEMORY_HOME_DIR_NAME,
331                        WRAP_STAGING_SUBDIR
332                    );
333                    tempfile::NamedTempFile::new()
334                        .context("ai-memory wrap: failed to create system-message tempfile")?
335                }
336            };
337            // Belt-and-braces: the tempfile crate already creates
338            // 0600 on Unix; pin it explicitly so a future tempfile
339            // upgrade can't silently loosen the boot-context file.
340            #[cfg(unix)]
341            {
342                use std::os::unix::fs::PermissionsExt;
343                let _ = std::fs::set_permissions(tf.path(), std::fs::Permissions::from_mode(0o600));
344            }
345            tf.write_all(system_msg.as_bytes())
346                .context("ai-memory wrap: failed to write system-message tempfile")?;
347            // Flush so the agent process reads the full message even
348            // if the OS hasn't drained the buffer yet.
349            tf.flush()
350                .context("ai-memory wrap: failed to flush system-message tempfile")?;
351            cmd.arg(flag).arg(tf.path().as_os_str());
352            for t in trailing {
353                cmd.arg(t);
354            }
355            tempfile_handle = Some(tf);
356        }
357        WrapStrategy::Auto => {
358            // Resolve and recurse. `Auto` should be handled by
359            // `resolve_strategy` before we get here, but if a caller
360            // synthesises a `WrapArgs` programmatically and leaves
361            // strategy as `Auto`, fall through to the lookup table.
362            let resolved = default_strategy(agent);
363            return build_command_for_strategy(agent, &resolved, system_msg, trailing);
364        }
365    }
366    Ok((cmd, tempfile_handle))
367}
368
369/// `ai-memory wrap` entry point. Returns the wrapped agent's exit code
370/// so `daemon_runtime` can `std::process::exit(code)` on a non-zero
371/// outcome — that's how shell pipelines and CI gates branch on the
372/// agent's success.
373///
374/// # Errors
375///
376/// - The wrapped agent binary cannot be spawned (`Command::status`
377///   surfaces the OS-level error).
378/// - `tempfile::NamedTempFile::new()` fails when the strategy is
379///   `MessageFile` (very rare; `/tmp` full or unwritable).
380pub fn run(
381    db_path: &Path,
382    args: &WrapArgs,
383    app_config: &crate::config::AppConfig,
384    _out: &mut CliOutput<'_>,
385) -> Result<i32> {
386    let strategy = resolve_strategy(args);
387
388    // Boot context. `--no-boot` skips it so the agent runs unwrapped
389    // (still through `Command::new(agent)` so this subcommand stays
390    // useful as a strategy-hooked launcher even with memory off).
391    let system_msg = if args.no_boot {
392        WRAP_PREAMBLE.to_string()
393    } else {
394        let boot_output = run_boot_capture(db_path, args.limit, args.budget_tokens, app_config);
395        build_system_message(&boot_output)
396    };
397
398    let (cmd, _tempfile_handle) =
399        build_command_for_strategy(&args.agent, &strategy, &system_msg, &args.trailing)?;
400
401    // _tempfile_handle is held by the local binding so it lives until
402    // after `spawn_and_wait` returns. Don't shorten its scope.
403    let code = spawn_and_wait(cmd)?;
404    Ok(code)
405}
406
407/// Public helper for callers (tests + future PR-7 recipe additions)
408/// that want to format an `OsStr` argv element back to UTF-8 for
409/// assertions / logging. Falls back to the lossy form so platforms
410/// with non-UTF-8 paths don't panic.
411#[must_use]
412pub fn os_str_to_string_lossy(s: &OsStr) -> String {
413    s.to_string_lossy().into_owned()
414}
415
416#[cfg(test)]
417mod tests {
418    use super::*;
419    use crate::cli::test_utils::{TestEnv, seed_memory};
420
421    fn default_args(agent: &str) -> WrapArgs {
422        WrapArgs {
423            agent: agent.to_string(),
424            system_flag: None,
425            system_env: None,
426            message_file_flag: None,
427            no_boot: false,
428            limit: DEFAULT_WRAP_LIMIT,
429            budget_tokens: DEFAULT_WRAP_BUDGET_TOKENS,
430            trailing: Vec::new(),
431        }
432    }
433
434    // NOTE: The canonical per-agent table pin moved to
435    // `crate::llm_cli_wrap::tests::default_strategy_per_known_agent_pins_1183`
436    // alongside the table itself in #1183. The tests below exercise the
437    // wrap-side dispatch (override precedence + command-build shape) and
438    // reach the moved table via the re-imported `default_strategy`
439    // symbol.
440
441    #[test]
442    fn resolve_strategy_explicit_overrides_lookup_table() {
443        let mut args = default_args("ollama");
444        args.system_flag = Some("--system-prompt".into());
445        // Even though "ollama" maps to SystemEnv in the lookup,
446        // explicit `--system-flag` wins.
447        assert_eq!(
448            resolve_strategy(&args),
449            WrapStrategy::SystemFlag {
450                flag: "--system-prompt".into()
451            }
452        );
453    }
454
455    #[test]
456    fn resolve_strategy_env_override_takes_precedence_over_flag_default() {
457        let mut args = default_args("codex");
458        args.system_env = Some("OPENAI_CLI_SYSTEM".into());
459        assert_eq!(
460            resolve_strategy(&args),
461            WrapStrategy::SystemEnv {
462                name: "OPENAI_CLI_SYSTEM".into()
463            }
464        );
465    }
466
467    #[test]
468    fn resolve_strategy_message_file_override() {
469        let mut args = default_args("codex");
470        args.message_file_flag = Some("--prompt-file".into());
471        assert_eq!(
472            resolve_strategy(&args),
473            WrapStrategy::MessageFile {
474                flag: "--prompt-file".into()
475            }
476        );
477    }
478
479    #[test]
480    fn build_system_message_prepends_preamble() {
481        let msg = build_system_message("- [mid/abc] hello");
482        assert!(msg.starts_with(WRAP_PREAMBLE));
483        assert!(msg.contains("hello"));
484        assert!(msg.contains("\n\n"), "preamble + body separator missing");
485    }
486
487    #[test]
488    fn build_system_message_empty_body_returns_preamble_only() {
489        let msg = build_system_message("");
490        assert_eq!(msg, WRAP_PREAMBLE);
491    }
492
493    #[test]
494    fn build_system_message_strips_trailing_whitespace() {
495        let msg = build_system_message("body line\n\n\n");
496        assert!(msg.ends_with("body line"));
497    }
498
499    #[test]
500    fn build_command_system_flag_sets_argv_correctly() {
501        let strat = WrapStrategy::SystemFlag {
502            flag: "--system".into(),
503        };
504        let trailing = vec![
505            "chat".to_string(),
506            "--model".to_string(),
507            "gpt-5".to_string(),
508        ];
509        let (cmd, tf) =
510            build_command_for_strategy("codex", &strat, "SYS-MSG-VALUE", &trailing).unwrap();
511        assert!(tf.is_none(), "SystemFlag must not allocate a tempfile");
512        let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
513        assert_eq!(
514            argv,
515            vec!["--system", "SYS-MSG-VALUE", "chat", "--model", "gpt-5"]
516        );
517        // Verify the program name (first arg of Command, not in
518        // get_args) — get_program is part of the std API.
519        assert_eq!(cmd.get_program(), OsStr::new("codex"));
520    }
521
522    #[test]
523    fn build_command_system_env_sets_env_var_and_omits_flag() {
524        let strat = WrapStrategy::SystemEnv {
525            name: "OLLAMA_SYSTEM".into(),
526        };
527        let trailing = vec!["run".to_string(), "hermes3:8b".to_string()];
528        let (cmd, tf) =
529            build_command_for_strategy("ollama", &strat, "SYS-ENV-MSG", &trailing).unwrap();
530        assert!(tf.is_none(), "SystemEnv must not allocate a tempfile");
531        let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
532        // The env-var strategy never injects a flag — argv is just the
533        // trailing args.
534        assert_eq!(argv, vec!["run", "hermes3:8b"]);
535        // Confirm OLLAMA_SYSTEM is set on the Command's env. get_envs()
536        // yields (key, Option<value>) pairs.
537        let env_pairs: Vec<(String, Option<String>)> = cmd
538            .get_envs()
539            .map(|(k, v)| {
540                (
541                    os_str_to_string_lossy(k),
542                    v.map(|x| os_str_to_string_lossy(x)),
543                )
544            })
545            .collect();
546        let entry = env_pairs
547            .iter()
548            .find(|(k, _)| k == "OLLAMA_SYSTEM")
549            .expect("OLLAMA_SYSTEM must be set");
550        assert_eq!(entry.1.as_deref(), Some("SYS-ENV-MSG"));
551    }
552
553    #[test]
554    fn wrap_strategy_message_file_creates_tempfile_and_cleans_up() {
555        let strat = WrapStrategy::MessageFile {
556            flag: "--message-file".into(),
557        };
558        let (path_owned, exists_during) = {
559            let (cmd, tf) =
560                build_command_for_strategy("aider", &strat, "FILE-MSG-CONTENT", &[]).unwrap();
561            let tf = tf.expect("MessageFile must allocate a tempfile");
562            // The argv should point at the tempfile path. We can't
563            // directly assert path equality on Windows (canonicalisation
564            // differs), so just check the `--message-file` flag is the
565            // first arg and the second arg is some non-empty path.
566            let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
567            assert_eq!(argv.len(), 2);
568            assert_eq!(argv[0], "--message-file");
569            assert!(!argv[1].is_empty());
570            // Sanity: the tempfile contains the expected message body.
571            let read_back = std::fs::read_to_string(tf.path()).unwrap();
572            assert_eq!(read_back, "FILE-MSG-CONTENT");
573            let exists = tf.path().exists();
574            // Take the path as PathBuf BEFORE dropping `tf` so we can
575            // re-stat after the block exits.
576            let p = tf.path().to_path_buf();
577            (p, exists)
578        };
579        assert!(
580            exists_during,
581            "tempfile must exist while NamedTempFile is alive"
582        );
583        // After the block ends, NamedTempFile is dropped, which
584        // unlinks the file (Unix and Windows both — tempfile crate
585        // smooths over the platform difference).
586        assert!(
587            !path_owned.exists(),
588            "tempfile must be cleaned up on Drop, but {} still exists",
589            path_owned.display()
590        );
591    }
592
593    /// #1575 — the boot-context message file must be staged under the
594    /// per-user ai-memory data dir (`~/.ai-memory/wrap/`, 0700 dir /
595    /// 0600 file), NOT the platform temp dir. The temp dir is only the
596    /// home-unresolvable fallback (exercised implicitly when
597    /// `dirs::home_dir()` returns `None`, which cannot be forced here
598    /// without unsafe env mutation — the fallback arm is plain
599    /// pre-#1575 behavior).
600    #[test]
601    fn message_file_staged_under_ai_memory_home_with_tight_perms_1575() {
602        let Some(staging) = message_file_staging_dir() else {
603            // No resolvable home in this environment — the WARN +
604            // temp-dir fallback arm applies; nothing to assert.
605            return;
606        };
607        let strat = WrapStrategy::MessageFile {
608            flag: "--message-file".into(),
609        };
610        let (_cmd, tf) =
611            build_command_for_strategy("aider", &strat, "BOOT-CONTEXT-1575", &[]).unwrap();
612        let tf = tf.expect("MessageFile must allocate a staged file");
613        assert_eq!(
614            tf.path().parent(),
615            Some(staging.as_path()),
616            "boot-context file must live under the ai-memory staging dir, got {}",
617            tf.path().display()
618        );
619        #[cfg(unix)]
620        {
621            use std::os::unix::fs::PermissionsExt;
622            let dmode = std::fs::metadata(&staging).unwrap().permissions().mode() & 0o777;
623            assert_eq!(dmode, 0o700, "staging dir must be 0700");
624            let fmode = std::fs::metadata(tf.path()).unwrap().permissions().mode() & 0o777;
625            assert_eq!(fmode, 0o600, "boot-context file must be 0600");
626        }
627    }
628
629    #[test]
630    fn wrap_with_unreachable_db_does_not_block_agent() {
631        // Boot honors `--quiet` and exits 0 with a warn header on stdout
632        // when the DB is missing. The captured stdout becomes the body
633        // of the wrap system message. We assert: (a) `run_boot_capture`
634        // returns *something* (the warn header) without erroring, and
635        // (b) the assembled system message still carries the preamble
636        // so the agent knows it has memory access (even if empty).
637        let env = TestEnv::fresh();
638        let bad = env
639            .db_path
640            .parent()
641            .unwrap()
642            .join("nope/that/does/not/exist/db.sqlite");
643        let captured = run_boot_capture(
644            &bad,
645            10,
646            DEFAULT_WRAP_BUDGET_TOKENS,
647            &crate::config::AppConfig::default(),
648        );
649        assert!(
650            captured.contains("# ai-memory boot: warn"),
651            "wrap should surface the warn header even with unreachable DB: {captured}"
652        );
653        let assembled = build_system_message(&captured);
654        assert!(assembled.starts_with(WRAP_PREAMBLE));
655        assert!(assembled.contains("warn"));
656    }
657
658    #[test]
659    fn wrap_with_no_boot_skips_context() {
660        // Smoke: the run path with `no_boot = true` produces a system
661        // message that's exactly the preamble (no boot body). We verify
662        // by re-running the equivalent assembly the `run` function uses
663        // when `args.no_boot` is true.
664        let mut args = default_args("codex");
665        args.no_boot = true;
666        // The `run` body's `if args.no_boot { WRAP_PREAMBLE.to_string() }`
667        // branch is what produces the system message in this mode.
668        // We replicate it here so we can assert on the value without
669        // spawning a subprocess (the real `codex` isn't on the test
670        // host's PATH).
671        let system_msg = if args.no_boot {
672            WRAP_PREAMBLE.to_string()
673        } else {
674            unreachable!()
675        };
676        assert_eq!(system_msg, WRAP_PREAMBLE);
677        // And the assembled command for that message must contain
678        // exactly the preamble as the flag value, no boot context.
679        let (cmd, _tf) = build_command_for_strategy(
680            &args.agent,
681            &resolve_strategy(&args),
682            &system_msg,
683            &args.trailing,
684        )
685        .unwrap();
686        let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
687        assert_eq!(argv.len(), 2);
688        assert_eq!(argv[0], "--system");
689        assert_eq!(argv[1], WRAP_PREAMBLE);
690    }
691
692    #[test]
693    fn wrap_injects_system_message_via_flag() {
694        // Seed a memory so the boot output is non-empty, then assert
695        // the assembled system message that wrap would pass to the
696        // agent contains both the preamble AND the seeded memory's
697        // title. This is the contract the docs/integrations recipes
698        // depend on.
699        let env = TestEnv::fresh();
700        seed_memory(&env.db_path, "ns-wrap-test", "wrap-injection-canary", "x");
701        let captured = run_boot_capture(
702            &env.db_path,
703            10,
704            DEFAULT_WRAP_BUDGET_TOKENS,
705            &crate::config::AppConfig::default(),
706        );
707        // boot::run sets the namespace from auto_namespace, which won't
708        // match `ns-wrap-test` unless cwd is set. The fallback path
709        // should still surface SOMETHING so the captured body is
710        // non-empty (warn or info header at minimum).
711        assert!(
712            !captured.is_empty(),
713            "expected non-empty boot capture, got empty"
714        );
715        let assembled = build_system_message(&captured);
716        assert!(assembled.starts_with(WRAP_PREAMBLE));
717        assert!(assembled.len() > WRAP_PREAMBLE.len());
718        // Now assert the assembled message rides through to the
719        // command's argv.
720        let (cmd, _tf) = build_command_for_strategy(
721            "codex",
722            &WrapStrategy::SystemFlag {
723                flag: "--system".into(),
724            },
725            &assembled,
726            &[],
727        )
728        .unwrap();
729        let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
730        assert_eq!(argv.len(), 2);
731        assert_eq!(argv[0], "--system");
732        assert!(argv[1].starts_with(WRAP_PREAMBLE));
733    }
734
735    #[test]
736    fn wrap_passes_through_exit_code_via_status_propagation() {
737        // We can't assume any specific binary is on PATH, but we can
738        // exercise the propagation logic with a guaranteed-available
739        // command: `false` on Unix exits 1, `true` exits 0. On Windows
740        // we use `cmd /C exit N`.
741        #[cfg(unix)]
742        {
743            // Exit 0
744            let cmd = Command::new("true");
745            let code = spawn_and_wait(cmd).unwrap();
746            assert_eq!(code, 0);
747            // Exit 1
748            let cmd = Command::new("false");
749            let code = spawn_and_wait(cmd).unwrap();
750            assert_eq!(code, 1);
751        }
752        #[cfg(windows)]
753        {
754            let mut cmd = Command::new("cmd");
755            cmd.args(["/C", "exit", "0"]);
756            let code = spawn_and_wait(cmd).unwrap();
757            assert_eq!(code, 0);
758            let mut cmd = Command::new("cmd");
759            cmd.args(["/C", "exit", "7"]);
760            let code = spawn_and_wait(cmd).unwrap();
761            assert_eq!(code, 7);
762        }
763    }
764
765    #[test]
766    fn wrap_run_returns_exit_code_for_real_subprocess() {
767        // End-to-end: drive `run` itself (not just the helpers). We
768        // wrap a known-good binary (`true` on unix, `cmd /C exit` on
769        // windows) and assert the returned code matches.
770        let mut env = TestEnv::fresh();
771        let db_path = env.db_path.clone();
772        let mut out = env.output();
773        #[cfg(unix)]
774        {
775            let mut args = default_args("true");
776            // Skip boot to avoid touching the DB and to keep the test
777            // deterministic. `--system "..."` is still passed to the
778            // agent — `true` ignores all argv, exits 0.
779            args.no_boot = true;
780            let code = run(
781                &db_path,
782                &args,
783                &crate::config::AppConfig::default(),
784                &mut out,
785            )
786            .unwrap();
787            assert_eq!(code, 0);
788        }
789        #[cfg(windows)]
790        {
791            let mut args = default_args("cmd");
792            args.no_boot = true;
793            // We override the strategy to SystemFlag with a no-op flag
794            // that `cmd /C` will ignore alongside the system message,
795            // then a real /C exit. Easier: override via system_env so
796            // no flag is added, then trailing carries `/C exit 5`.
797            args.system_env = Some("WRAP_DUMMY".into());
798            args.trailing = vec!["/C".into(), "exit".into(), "5".into()];
799            let code = run(
800                &db_path,
801                &args,
802                &crate::config::AppConfig::default(),
803                &mut out,
804            )
805            .unwrap();
806            assert_eq!(code, 5);
807        }
808    }
809
810    #[test]
811    fn auto_strategy_resolves_at_command_build_time() {
812        // Exercise the `WrapStrategy::Auto` recursive branch in
813        // `build_command_for_strategy`.
814        let (cmd, tf) = build_command_for_strategy(
815            "codex",
816            &WrapStrategy::Auto,
817            "AUTO-MSG",
818            &["chat".to_string()],
819        )
820        .unwrap();
821        assert!(tf.is_none());
822        let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
823        // codex auto-resolves to SystemFlag{--system}.
824        assert_eq!(argv, vec!["--system", "AUTO-MSG", "chat"]);
825    }
826
827    #[test]
828    fn auto_strategy_resolves_to_message_file_for_aider() {
829        let (cmd, tf) =
830            build_command_for_strategy("aider", &WrapStrategy::Auto, "AIDER-MSG", &[]).unwrap();
831        // aider auto-resolves to MessageFile, so a tempfile must be
832        // allocated.
833        assert!(tf.is_some());
834        let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
835        assert_eq!(argv.len(), 2);
836        assert_eq!(argv[0], "--message-file");
837    }
838
839    #[test]
840    fn run_boot_capture_returns_string_not_panics_on_missing_db() {
841        // Hardening: every error path inside boot must surface as a
842        // String (possibly empty, possibly the warn header) — never a
843        // panic — so the wrapped agent always runs.
844        let env = TestEnv::fresh();
845        let bad = env
846            .db_path
847            .parent()
848            .unwrap()
849            .join("__definitely_missing__/db");
850        let s = run_boot_capture(
851            &bad,
852            10,
853            DEFAULT_WRAP_BUDGET_TOKENS,
854            &crate::config::AppConfig::default(),
855        );
856        // Either the warn header or empty (both are non-panic outcomes).
857        assert!(
858            s.is_empty() || s.contains("# ai-memory boot:"),
859            "expected warn header or empty, got: {s}"
860        );
861    }
862
863    /// Coverage restoration (post-#1575 floor dip): the
864    /// `boot::run(...).is_err()` hard-failure arm in
865    /// `run_boot_capture` must return an EMPTY string (agent runs
866    /// unwrapped) — forced by pointing db_path at a DIRECTORY, which
867    /// the sqlite open cannot create-or-open even under `--quiet`.
868    #[test]
869    fn run_boot_capture_returns_empty_when_db_path_is_a_directory() {
870        let env = TestEnv::fresh();
871        let dir_as_db = env.db_path.parent().unwrap().to_path_buf();
872        let s = run_boot_capture(
873            &dir_as_db,
874            10,
875            DEFAULT_WRAP_BUDGET_TOKENS,
876            &crate::config::AppConfig::default(),
877        );
878        assert!(
879            s.is_empty() || s.contains("# ai-memory boot:"),
880            "directory-as-db must yield empty or warn-header output, got: {s}"
881        );
882    }
883
884    /// Coverage restoration: the MessageFile arm's trailing-arg
885    /// passthrough loop — trailing CLI args must land on the wrapped
886    /// command AFTER the message-file flag pair.
887    #[test]
888    fn message_file_strategy_passes_trailing_args_through() {
889        let strat = WrapStrategy::MessageFile {
890            flag: "--message-file".into(),
891        };
892        let trailing = vec!["--model".to_string(), "gpt-x".to_string()];
893        let (cmd, tf) =
894            build_command_for_strategy("aider", &strat, "BOOT-TRAIL", &trailing).unwrap();
895        let _tf = tf.expect("MessageFile must allocate a staged file");
896        let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
897        assert_eq!(argv[0], "--message-file");
898        assert_eq!(
899            &argv[2..],
900            ["--model", "gpt-x"],
901            "trailing args must follow the message-file pair: {argv:?}"
902        );
903    }
904}