use std::sync::Mutex;
use std::time::{SystemTime, UNIX_EPOCH};
use anyhow::{Context, Result};
use super::config::resolve_model_for_api;
use super::pricing::{estimate_cost_usd, pricing_for_model};
use super::stable_working_dir;
use super::TokenUsage;
static ENV_LOCK: Mutex<()> = Mutex::new(());
fn with_env_vars<T>(vars: &[(&str, Option<&str>)], f: impl FnOnce() -> T) -> T {
let _guard = ENV_LOCK.lock().expect("env lock should acquire");
let old_values = vars
.iter()
.map(|(key, _)| ((*key).to_string(), std::env::var(key).ok()))
.collect::<Vec<_>>();
for (key, value) in vars {
match value {
Some(value) => unsafe { std::env::set_var(key, value) },
None => unsafe { std::env::remove_var(key) },
}
}
let result = f();
for (key, value) in old_values {
match value {
Some(value) => unsafe { std::env::set_var(&key, value) },
None => unsafe { std::env::remove_var(&key) },
}
}
result
}
#[test]
fn resolve_model_for_api_maps_short_names() {
assert_eq!(resolve_model_for_api("haiku"), "claude-haiku-4-5-20251001");
assert_eq!(
resolve_model_for_api("sonnet"),
"claude-sonnet-4-5-20250514"
);
assert_eq!(resolve_model_for_api("opus"), "claude-opus-4-20250514");
assert_eq!(resolve_model_for_api("custom-model"), "custom-model");
}
#[test]
fn pricing_for_model_uses_model_defaults() {
with_env_vars(
&[
("REMEM_PRICE_INPUT_PER_MTOK", None),
("REMEM_PRICE_OUTPUT_PER_MTOK", None),
("REMEM_PRICE_HAIKU_INPUT_PER_MTOK", None),
("REMEM_PRICE_HAIKU_OUTPUT_PER_MTOK", None),
],
|| {
assert_eq!(pricing_for_model("haiku"), (1.0, 5.0));
},
);
}
#[test]
fn pricing_for_gpt_52_uses_current_flagship_rate() {
with_env_vars(
&[
("REMEM_PRICE_INPUT_PER_MTOK", None),
("REMEM_PRICE_OUTPUT_PER_MTOK", None),
("REMEM_PRICE_GPT5_CODEX_INPUT_PER_MTOK", None),
("REMEM_PRICE_GPT5_CODEX_OUTPUT_PER_MTOK", None),
],
|| {
assert_eq!(pricing_for_model("gpt-5.2"), (1.75, 14.0));
},
);
}
#[test]
fn pricing_for_model_prefers_env_override() {
with_env_vars(
&[
("REMEM_PRICE_INPUT_PER_MTOK", Some("1.25")),
("REMEM_PRICE_OUTPUT_PER_MTOK", Some("6.5")),
],
|| {
assert_eq!(pricing_for_model("sonnet"), (1.25, 6.5));
},
);
}
#[test]
fn estimate_cost_usd_combines_input_and_output_prices() {
with_env_vars(
&[
("REMEM_PRICE_INPUT_PER_MTOK", Some("2.0")),
("REMEM_PRICE_OUTPUT_PER_MTOK", Some("8.0")),
],
|| {
let usage = TokenUsage::estimated(500_000, 250_000);
let (cost, pricing_source) = estimate_cost_usd("any-model", &usage);
assert_eq!(pricing_source, "env_override");
assert!((cost - 3.0).abs() < f64::EPSILON);
},
);
}
#[test]
fn estimate_cost_usd_charges_cache_and_reasoning_separately() {
with_env_vars(
&[
("REMEM_PRICE_INPUT_PER_MTOK", None),
("REMEM_PRICE_OUTPUT_PER_MTOK", None),
("REMEM_PRICE_REASONING_PER_MTOK", None),
("REMEM_PRICE_CACHE_READ_PER_MTOK", None),
("REMEM_PRICE_CACHE_CREATION_PER_MTOK", None),
],
|| {
let usage = TokenUsage {
input_tokens: 1_000_000,
output_tokens: 1_000_000,
reasoning_tokens: 1_000_000,
cache_read_tokens: 1_000_000,
..TokenUsage::default()
};
let (cost, pricing_source) = estimate_cost_usd("gpt-5.5", &usage);
assert_eq!(pricing_source, "remem_static");
assert!((cost - 65.5).abs() < f64::EPSILON);
},
);
}
#[test]
fn stable_working_dir_uses_data_dir_even_if_caller_cwd_disappears() {
let data_dir = crate::db::test_support::ScopedTestDataDir::new("ai-stable-cwd");
let got = stable_working_dir();
assert_eq!(got, data_dir.path);
assert!(got.is_dir());
}
#[cfg(unix)]
#[test]
fn claude_cli_child_disables_remem_hooks() -> Result<()> {
use std::os::unix::fs::PermissionsExt;
let unique = SystemTime::now()
.duration_since(UNIX_EPOCH)
.context("system time should be after unix epoch")?
.as_nanos();
let temp_dir = std::env::temp_dir().join(format!(
"remem-fake-claude-{}-{}",
std::process::id(),
unique
));
std::fs::create_dir_all(&temp_dir).context("fake claude temp dir should be created")?;
let script_path = temp_dir.join("claude");
let env_path = temp_dir.join("env.txt");
std::fs::write(
&script_path,
r#"#!/bin/sh
{
printf 'REMEM_DISABLE_HOOKS=%s\n' "${REMEM_DISABLE_HOOKS-}"
printf 'CLAUDECODE=%s\n' "${CLAUDECODE-__unset__}"
} > "$REMEM_FAKE_CLAUDE_ENV_OUT"
cat >/dev/null
printf 'ok\n'
"#,
)
.context("fake claude script should be written")?;
let mut permissions = std::fs::metadata(&script_path)
.context("fake claude metadata should be readable")?
.permissions();
permissions.set_mode(0o700);
std::fs::set_permissions(&script_path, permissions)
.context("fake claude script should be executable")?;
let script = script_path
.to_str()
.context("fake claude path should be valid utf-8")?
.to_string();
let env_out = env_path
.to_str()
.context("fake claude env path should be valid utf-8")?
.to_string();
with_env_vars(
&[
("REMEM_FAKE_CLAUDE_ENV_OUT", Some(env_out.as_str())),
("REMEM_DISABLE_HOOKS", None),
("CLAUDECODE", Some("parent-session")),
],
|| -> Result<()> {
let profile = crate::runtime_config::ResolvedMemoryAiProfile {
profile_name: "test-claude".to_string(),
executor: crate::runtime_config::MemoryAiExecutor::ClaudeCli,
model: Some("haiku".to_string()),
cli_path: Some(script.clone()),
base_url: None,
reasoning_effort: None,
};
let _data_dir = crate::db::test_support::ScopedTestDataDir::new("ai-claude-child");
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.context("tokio runtime should build")?;
let result = runtime
.block_on(super::cli::call_cli("system", "user", &profile))
.context("fake claude call should succeed")?;
assert_eq!(result.text, "ok");
Ok(())
},
)?;
let captured =
std::fs::read_to_string(&env_path).context("fake claude should capture child env")?;
assert!(captured.contains("REMEM_DISABLE_HOOKS=1"), "{captured:?}");
assert!(captured.contains("CLAUDECODE=__unset__"), "{captured:?}");
std::fs::remove_dir_all(&temp_dir)
.with_context(|| format!("failed to remove {}", temp_dir.display()))?;
Ok(())
}
#[cfg(unix)]
#[tokio::test(flavor = "current_thread")]
async fn resolved_profile_scope_does_not_reread_removed_config_profile() -> Result<()> {
use std::os::unix::fs::PermissionsExt;
let unique = SystemTime::now()
.duration_since(UNIX_EPOCH)
.context("system time should be after unix epoch")?
.as_nanos();
let temp_dir = std::env::temp_dir().join(format!(
"remem-frozen-profile-{}-{unique}",
std::process::id()
));
std::fs::create_dir_all(&temp_dir)?;
let script_path = temp_dir.join("claude");
std::fs::write(
&script_path,
"#!/bin/sh\ncat >/dev/null\nprintf 'frozen-profile\\n'\n",
)?;
let mut permissions = std::fs::metadata(&script_path)?.permissions();
permissions.set_mode(0o700);
std::fs::set_permissions(&script_path, permissions)?;
let _data_dir = crate::db::test_support::ScopedTestDataDir::new("ai-frozen-profile");
let profile = crate::runtime_config::ResolvedMemoryAiProfile {
profile_name: "removed-after-validation".to_string(),
executor: crate::runtime_config::MemoryAiExecutor::ClaudeCli,
model: Some("haiku".to_string()),
cli_path: Some(script_path.to_string_lossy().into_owned()),
base_url: None,
reasoning_effort: None,
};
let response = super::with_resolved_profile(
profile,
super::call_ai(
"system",
"user",
super::UsageContext {
project: Some("/tmp/remem"),
session_id: Some("frozen-profile-session"),
operation: "frozen_profile_test",
host: None,
profile: Some("removed-after-validation"),
},
),
)
.await?;
assert_eq!(response, "frozen-profile");
std::fs::remove_dir_all(&temp_dir)?;
Ok(())
}