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}