use crate::cli::CliOutput;
use crate::cli::boot::{self, BootArgs};
use crate::llm_cli_wrap::{WrapStrategy, default_strategy};
use anyhow::{Context, Result};
use clap::Args;
use std::ffi::OsStr;
use std::io::Write;
use std::path::Path;
use std::process::{Command, Stdio};
const DEFAULT_WRAP_BUDGET_TOKENS: usize = 4096;
const DEFAULT_WRAP_LIMIT: usize = 10;
const WRAP_PREAMBLE: &str = "You have access to ai-memory, a persistent memory system. \
The recent context loaded for you appears below. Reference it when relevant to the user's request.";
const WRAP_STAGING_SUBDIR: &str = "wrap";
fn message_file_staging_dir() -> Option<std::path::PathBuf> {
let dir = dirs::home_dir()?
.join(crate::AI_MEMORY_HOME_DIR_NAME)
.join(WRAP_STAGING_SUBDIR);
std::fs::create_dir_all(&dir).ok()?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o700)).ok()?;
}
Some(dir)
}
#[derive(Args, Debug)]
pub struct WrapArgs {
pub agent: String,
#[arg(long, value_name = "FLAG")]
pub system_flag: Option<String>,
#[arg(long, value_name = "NAME", conflicts_with_all = ["system_flag", "message_file_flag"])]
pub system_env: Option<String>,
#[arg(long, value_name = "FLAG", conflicts_with_all = ["system_flag", "system_env"])]
pub message_file_flag: Option<String>,
#[arg(long, default_value_t = false)]
pub no_boot: bool,
#[arg(long, default_value_t = DEFAULT_WRAP_LIMIT)]
pub limit: usize,
#[arg(long, default_value_t = DEFAULT_WRAP_BUDGET_TOKENS)]
pub budget_tokens: usize,
#[arg(trailing_var_arg = true, allow_hyphen_values = true)]
pub trailing: Vec<String>,
}
fn resolve_strategy(args: &WrapArgs) -> WrapStrategy {
if let Some(name) = args.system_env.as_deref() {
return WrapStrategy::SystemEnv { name: name.into() };
}
if let Some(flag) = args.message_file_flag.as_deref() {
return WrapStrategy::MessageFile { flag: flag.into() };
}
if let Some(flag) = args.system_flag.as_deref() {
return WrapStrategy::SystemFlag { flag: flag.into() };
}
default_strategy(&args.agent)
}
fn run_boot_capture(
db_path: &Path,
limit: usize,
budget_tokens: usize,
app_config: &crate::config::AppConfig,
) -> String {
let mut stdout: Vec<u8> = Vec::new();
let mut stderr: Vec<u8> = Vec::new();
let mut out = CliOutput::from_std(&mut stdout, &mut stderr);
let args = BootArgs {
namespace: None,
limit,
budget_tokens,
format: "text".to_string(),
no_header: false,
quiet: true,
cwd: None,
};
if boot::run(db_path, &args, app_config, &mut out).is_err() {
return String::new();
}
String::from_utf8(stdout).unwrap_or_default()
}
fn build_system_message(boot_output: &str) -> String {
let trimmed = boot_output.trim_end();
if trimmed.is_empty() {
WRAP_PREAMBLE.to_string()
} else {
format!("{WRAP_PREAMBLE}\n\n{trimmed}")
}
}
fn spawn_and_wait(mut cmd: Command) -> Result<i32> {
cmd.stdin(Stdio::inherit())
.stdout(Stdio::inherit())
.stderr(Stdio::inherit());
let status = cmd
.status()
.with_context(|| format!("ai-memory wrap: failed to spawn agent {cmd:?}"))?;
let code = if let Some(c) = status.code() {
c
} else {
#[cfg(unix)]
{
use std::os::unix::process::ExitStatusExt;
status.signal().map_or(1, |s| 128 + s)
}
#[cfg(not(unix))]
{
1
}
};
Ok(code)
}
fn build_command_for_strategy(
agent: &str,
strategy: &WrapStrategy,
system_msg: &str,
trailing: &[String],
) -> Result<(Command, Option<tempfile::NamedTempFile>)> {
let mut cmd = Command::new(agent);
let mut tempfile_handle: Option<tempfile::NamedTempFile> = None;
match strategy {
WrapStrategy::SystemFlag { flag } => {
cmd.arg(flag).arg(system_msg);
for t in trailing {
cmd.arg(t);
}
}
WrapStrategy::SystemEnv { name } => {
cmd.env(name, system_msg);
for t in trailing {
cmd.arg(t);
}
}
WrapStrategy::MessageFile { flag } => {
let mut tf = match message_file_staging_dir() {
Some(dir) => tempfile::NamedTempFile::new_in(&dir).context(
"ai-memory wrap: failed to create system-message file in staging dir",
)?,
None => {
tracing::warn!(
"ai-memory wrap: could not resolve/secure the {}/{} staging dir under \
the home directory; falling back to the platform temp dir for the \
boot-context message file (#1575 — memory contents will transit a \
shared temp path)",
crate::AI_MEMORY_HOME_DIR_NAME,
WRAP_STAGING_SUBDIR
);
tempfile::NamedTempFile::new()
.context("ai-memory wrap: failed to create system-message tempfile")?
}
};
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let _ = std::fs::set_permissions(tf.path(), std::fs::Permissions::from_mode(0o600));
}
tf.write_all(system_msg.as_bytes())
.context("ai-memory wrap: failed to write system-message tempfile")?;
tf.flush()
.context("ai-memory wrap: failed to flush system-message tempfile")?;
cmd.arg(flag).arg(tf.path().as_os_str());
for t in trailing {
cmd.arg(t);
}
tempfile_handle = Some(tf);
}
WrapStrategy::Auto => {
let resolved = default_strategy(agent);
return build_command_for_strategy(agent, &resolved, system_msg, trailing);
}
}
Ok((cmd, tempfile_handle))
}
pub fn run(
db_path: &Path,
args: &WrapArgs,
app_config: &crate::config::AppConfig,
_out: &mut CliOutput<'_>,
) -> Result<i32> {
let strategy = resolve_strategy(args);
let system_msg = if args.no_boot {
WRAP_PREAMBLE.to_string()
} else {
let boot_output = run_boot_capture(db_path, args.limit, args.budget_tokens, app_config);
build_system_message(&boot_output)
};
let (cmd, _tempfile_handle) =
build_command_for_strategy(&args.agent, &strategy, &system_msg, &args.trailing)?;
let code = spawn_and_wait(cmd)?;
Ok(code)
}
#[must_use]
pub fn os_str_to_string_lossy(s: &OsStr) -> String {
s.to_string_lossy().into_owned()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::cli::test_utils::{TestEnv, seed_memory};
fn default_args(agent: &str) -> WrapArgs {
WrapArgs {
agent: agent.to_string(),
system_flag: None,
system_env: None,
message_file_flag: None,
no_boot: false,
limit: DEFAULT_WRAP_LIMIT,
budget_tokens: DEFAULT_WRAP_BUDGET_TOKENS,
trailing: Vec::new(),
}
}
#[test]
fn resolve_strategy_explicit_overrides_lookup_table() {
let mut args = default_args("ollama");
args.system_flag = Some("--system-prompt".into());
assert_eq!(
resolve_strategy(&args),
WrapStrategy::SystemFlag {
flag: "--system-prompt".into()
}
);
}
#[test]
fn resolve_strategy_env_override_takes_precedence_over_flag_default() {
let mut args = default_args("codex");
args.system_env = Some("OPENAI_CLI_SYSTEM".into());
assert_eq!(
resolve_strategy(&args),
WrapStrategy::SystemEnv {
name: "OPENAI_CLI_SYSTEM".into()
}
);
}
#[test]
fn resolve_strategy_message_file_override() {
let mut args = default_args("codex");
args.message_file_flag = Some("--prompt-file".into());
assert_eq!(
resolve_strategy(&args),
WrapStrategy::MessageFile {
flag: "--prompt-file".into()
}
);
}
#[test]
fn build_system_message_prepends_preamble() {
let msg = build_system_message("- [mid/abc] hello");
assert!(msg.starts_with(WRAP_PREAMBLE));
assert!(msg.contains("hello"));
assert!(msg.contains("\n\n"), "preamble + body separator missing");
}
#[test]
fn build_system_message_empty_body_returns_preamble_only() {
let msg = build_system_message("");
assert_eq!(msg, WRAP_PREAMBLE);
}
#[test]
fn build_system_message_strips_trailing_whitespace() {
let msg = build_system_message("body line\n\n\n");
assert!(msg.ends_with("body line"));
}
#[test]
fn build_command_system_flag_sets_argv_correctly() {
let strat = WrapStrategy::SystemFlag {
flag: "--system".into(),
};
let trailing = vec![
"chat".to_string(),
"--model".to_string(),
"gpt-5".to_string(),
];
let (cmd, tf) =
build_command_for_strategy("codex", &strat, "SYS-MSG-VALUE", &trailing).unwrap();
assert!(tf.is_none(), "SystemFlag must not allocate a tempfile");
let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
assert_eq!(
argv,
vec!["--system", "SYS-MSG-VALUE", "chat", "--model", "gpt-5"]
);
assert_eq!(cmd.get_program(), OsStr::new("codex"));
}
#[test]
fn build_command_system_env_sets_env_var_and_omits_flag() {
let strat = WrapStrategy::SystemEnv {
name: "OLLAMA_SYSTEM".into(),
};
let trailing = vec!["run".to_string(), "hermes3:8b".to_string()];
let (cmd, tf) =
build_command_for_strategy("ollama", &strat, "SYS-ENV-MSG", &trailing).unwrap();
assert!(tf.is_none(), "SystemEnv must not allocate a tempfile");
let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
assert_eq!(argv, vec!["run", "hermes3:8b"]);
let env_pairs: Vec<(String, Option<String>)> = cmd
.get_envs()
.map(|(k, v)| {
(
os_str_to_string_lossy(k),
v.map(|x| os_str_to_string_lossy(x)),
)
})
.collect();
let entry = env_pairs
.iter()
.find(|(k, _)| k == "OLLAMA_SYSTEM")
.expect("OLLAMA_SYSTEM must be set");
assert_eq!(entry.1.as_deref(), Some("SYS-ENV-MSG"));
}
#[test]
fn wrap_strategy_message_file_creates_tempfile_and_cleans_up() {
let strat = WrapStrategy::MessageFile {
flag: "--message-file".into(),
};
let (path_owned, exists_during) = {
let (cmd, tf) =
build_command_for_strategy("aider", &strat, "FILE-MSG-CONTENT", &[]).unwrap();
let tf = tf.expect("MessageFile must allocate a tempfile");
let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
assert_eq!(argv.len(), 2);
assert_eq!(argv[0], "--message-file");
assert!(!argv[1].is_empty());
let read_back = std::fs::read_to_string(tf.path()).unwrap();
assert_eq!(read_back, "FILE-MSG-CONTENT");
let exists = tf.path().exists();
let p = tf.path().to_path_buf();
(p, exists)
};
assert!(
exists_during,
"tempfile must exist while NamedTempFile is alive"
);
assert!(
!path_owned.exists(),
"tempfile must be cleaned up on Drop, but {} still exists",
path_owned.display()
);
}
#[test]
fn message_file_staged_under_ai_memory_home_with_tight_perms_1575() {
let Some(staging) = message_file_staging_dir() else {
return;
};
let strat = WrapStrategy::MessageFile {
flag: "--message-file".into(),
};
let (_cmd, tf) =
build_command_for_strategy("aider", &strat, "BOOT-CONTEXT-1575", &[]).unwrap();
let tf = tf.expect("MessageFile must allocate a staged file");
assert_eq!(
tf.path().parent(),
Some(staging.as_path()),
"boot-context file must live under the ai-memory staging dir, got {}",
tf.path().display()
);
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let dmode = std::fs::metadata(&staging).unwrap().permissions().mode() & 0o777;
assert_eq!(dmode, 0o700, "staging dir must be 0700");
let fmode = std::fs::metadata(tf.path()).unwrap().permissions().mode() & 0o777;
assert_eq!(fmode, 0o600, "boot-context file must be 0600");
}
}
#[test]
fn wrap_with_unreachable_db_does_not_block_agent() {
let env = TestEnv::fresh();
let bad = env
.db_path
.parent()
.unwrap()
.join("nope/that/does/not/exist/db.sqlite");
let captured = run_boot_capture(
&bad,
10,
DEFAULT_WRAP_BUDGET_TOKENS,
&crate::config::AppConfig::default(),
);
assert!(
captured.contains("# ai-memory boot: warn"),
"wrap should surface the warn header even with unreachable DB: {captured}"
);
let assembled = build_system_message(&captured);
assert!(assembled.starts_with(WRAP_PREAMBLE));
assert!(assembled.contains("warn"));
}
#[test]
fn wrap_with_no_boot_skips_context() {
let mut args = default_args("codex");
args.no_boot = true;
let system_msg = if args.no_boot {
WRAP_PREAMBLE.to_string()
} else {
unreachable!()
};
assert_eq!(system_msg, WRAP_PREAMBLE);
let (cmd, _tf) = build_command_for_strategy(
&args.agent,
&resolve_strategy(&args),
&system_msg,
&args.trailing,
)
.unwrap();
let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
assert_eq!(argv.len(), 2);
assert_eq!(argv[0], "--system");
assert_eq!(argv[1], WRAP_PREAMBLE);
}
#[test]
fn wrap_injects_system_message_via_flag() {
let env = TestEnv::fresh();
seed_memory(&env.db_path, "ns-wrap-test", "wrap-injection-canary", "x");
let captured = run_boot_capture(
&env.db_path,
10,
DEFAULT_WRAP_BUDGET_TOKENS,
&crate::config::AppConfig::default(),
);
assert!(
!captured.is_empty(),
"expected non-empty boot capture, got empty"
);
let assembled = build_system_message(&captured);
assert!(assembled.starts_with(WRAP_PREAMBLE));
assert!(assembled.len() > WRAP_PREAMBLE.len());
let (cmd, _tf) = build_command_for_strategy(
"codex",
&WrapStrategy::SystemFlag {
flag: "--system".into(),
},
&assembled,
&[],
)
.unwrap();
let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
assert_eq!(argv.len(), 2);
assert_eq!(argv[0], "--system");
assert!(argv[1].starts_with(WRAP_PREAMBLE));
}
#[test]
fn wrap_passes_through_exit_code_via_status_propagation() {
#[cfg(unix)]
{
let cmd = Command::new("true");
let code = spawn_and_wait(cmd).unwrap();
assert_eq!(code, 0);
let cmd = Command::new("false");
let code = spawn_and_wait(cmd).unwrap();
assert_eq!(code, 1);
}
#[cfg(windows)]
{
let mut cmd = Command::new("cmd");
cmd.args(["/C", "exit", "0"]);
let code = spawn_and_wait(cmd).unwrap();
assert_eq!(code, 0);
let mut cmd = Command::new("cmd");
cmd.args(["/C", "exit", "7"]);
let code = spawn_and_wait(cmd).unwrap();
assert_eq!(code, 7);
}
}
#[test]
fn wrap_run_returns_exit_code_for_real_subprocess() {
let mut env = TestEnv::fresh();
let db_path = env.db_path.clone();
let mut out = env.output();
#[cfg(unix)]
{
let mut args = default_args("true");
args.no_boot = true;
let code = run(
&db_path,
&args,
&crate::config::AppConfig::default(),
&mut out,
)
.unwrap();
assert_eq!(code, 0);
}
#[cfg(windows)]
{
let mut args = default_args("cmd");
args.no_boot = true;
args.system_env = Some("WRAP_DUMMY".into());
args.trailing = vec!["/C".into(), "exit".into(), "5".into()];
let code = run(
&db_path,
&args,
&crate::config::AppConfig::default(),
&mut out,
)
.unwrap();
assert_eq!(code, 5);
}
}
#[test]
fn auto_strategy_resolves_at_command_build_time() {
let (cmd, tf) = build_command_for_strategy(
"codex",
&WrapStrategy::Auto,
"AUTO-MSG",
&["chat".to_string()],
)
.unwrap();
assert!(tf.is_none());
let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
assert_eq!(argv, vec!["--system", "AUTO-MSG", "chat"]);
}
#[test]
fn auto_strategy_resolves_to_message_file_for_aider() {
let (cmd, tf) =
build_command_for_strategy("aider", &WrapStrategy::Auto, "AIDER-MSG", &[]).unwrap();
assert!(tf.is_some());
let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
assert_eq!(argv.len(), 2);
assert_eq!(argv[0], "--message-file");
}
#[test]
fn run_boot_capture_returns_string_not_panics_on_missing_db() {
let env = TestEnv::fresh();
let bad = env
.db_path
.parent()
.unwrap()
.join("__definitely_missing__/db");
let s = run_boot_capture(
&bad,
10,
DEFAULT_WRAP_BUDGET_TOKENS,
&crate::config::AppConfig::default(),
);
assert!(
s.is_empty() || s.contains("# ai-memory boot:"),
"expected warn header or empty, got: {s}"
);
}
#[test]
fn run_boot_capture_returns_empty_when_db_path_is_a_directory() {
let env = TestEnv::fresh();
let dir_as_db = env.db_path.parent().unwrap().to_path_buf();
let s = run_boot_capture(
&dir_as_db,
10,
DEFAULT_WRAP_BUDGET_TOKENS,
&crate::config::AppConfig::default(),
);
assert!(
s.is_empty() || s.contains("# ai-memory boot:"),
"directory-as-db must yield empty or warn-header output, got: {s}"
);
}
#[test]
fn message_file_strategy_passes_trailing_args_through() {
let strat = WrapStrategy::MessageFile {
flag: "--message-file".into(),
};
let trailing = vec!["--model".to_string(), "gpt-x".to_string()];
let (cmd, tf) =
build_command_for_strategy("aider", &strat, "BOOT-TRAIL", &trailing).unwrap();
let _tf = tf.expect("MessageFile must allocate a staged file");
let argv: Vec<String> = cmd.get_args().map(|s| os_str_to_string_lossy(s)).collect();
assert_eq!(argv[0], "--message-file");
assert_eq!(
&argv[2..],
["--model", "gpt-x"],
"trailing args must follow the message-file pair: {argv:?}"
);
}
}