use std::collections::{BTreeMap, BTreeSet};
use std::path::{Path, PathBuf};
use config::{Config, Environment, File, FileFormat};
use serde_json::Value;
use super::{Settings, validate};
pub const CONFIG_FILE: &str = "overbrainer.toml";
pub const ENV_PREFIX: &str = "OVERBRAINER";
const NOT_CONFIG_VARIABLES: [&str; 4] = [
"OVERBRAINER_TUI_",
"OVERBRAINER_TEST_",
crate::update::NO_UPDATE_CHECK_ENV,
super::SSH_CLIENT_ENV,
];
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EnvSource {
Process,
Vars(Vec<(String, String)>),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Source {
pub text: Option<String>,
pub env: EnvSource,
}
impl Source {
pub fn load(&self, project_dir: &Path) -> Result<Settings, ConfigError> {
match &self.text {
Some(text) => load_str(text, self.env.clone()),
None => load(project_dir, self.env.clone()),
}
}
}
impl From<EnvSource> for Source {
fn from(env: EnvSource) -> Self {
Self { text: None, env }
}
}
#[derive(Debug, thiserror::Error)]
pub enum ConfigError {
#[error("cannot read {}: {error}", path.display())]
Read {
path: PathBuf,
error: std::io::Error,
},
#[error("invalid configuration: {0}")]
Parse(String),
#[error("invalid configuration:\n {}", .0.join("\n "))]
Invalid(Vec<String>),
}
impl From<config::ConfigError> for ConfigError {
fn from(source: config::ConfigError) -> Self {
Self::Parse(describe(&source))
}
}
fn describe(error: &config::ConfigError) -> String {
match error {
config::ConfigError::Type { key, expected, .. } => {
let key = key.as_deref().unwrap_or("a value");
format!("`{key}` has the wrong type, expected {expected}")
},
config::ConfigError::At { error, key, .. } => match key {
Some(key) => format!("`{key}`: {}", describe(error)),
None => describe(error),
},
config::ConfigError::Message(message) => redact(message),
config::ConfigError::FileParse { .. } => redact_toml_error(error),
_ => "invalid configuration".to_string(),
}
}
fn redact_toml_error(error: &config::ConfigError) -> String {
let message = error.to_string();
let first_line = message.lines().next().unwrap_or("");
match first_line.find("at line ") {
Some(pos) => format!(
"TOML syntax error in overbrainer.toml {}",
first_line[pos..].trim_end()
),
None => "TOML syntax error in overbrainer.toml".to_string(),
}
}
fn redact(message: &str) -> String {
if message.starts_with("unknown field") || message.starts_with("missing field") {
return message.to_string();
}
if let Some(index) = message.find("expected") {
return format!("invalid value, {}", &message[index..]);
}
if let Some((head, _)) = message.split_once(" does not have variant constructor") {
return format!("{head}: unknown variant");
}
"invalid value".to_string()
}
pub fn load(project_dir: &Path, env: EnvSource) -> Result<Settings, ConfigError> {
let path = project_dir.join(CONFIG_FILE);
let content = std::fs::read_to_string(&path).map_err(|error| ConfigError::Read {
path: path.clone(),
error,
})?;
load_str(&content, env)
}
pub fn load_str(content: &str, env: EnvSource) -> Result<Settings, ConfigError> {
let env: config::Map<String, String> = env_pairs(env)
.into_iter()
.filter(|(key, _)| !is_not_config(key))
.collect();
let file_only = Config::builder()
.add_source(File::from_str(content, FileFormat::Toml))
.build()?;
let mut problems = validate::env_only_in_file(&file_only);
let mut settings: Settings = Config::builder()
.add_source(File::from_str(content, FileFormat::Toml))
.add_source(
Environment::with_prefix(ENV_PREFIX)
.prefix_separator("_")
.separator("__")
.try_parsing(false)
.source(Some(env)),
)
.build()?
.try_deserialize()?;
if let Some(training) = settings.training.as_mut() {
let file_extra: BTreeMap<String, Value> =
file_only.get("training.axolotl_extra").unwrap_or_default();
coerce_env_scalars(&mut training.axolotl_extra, &file_extra);
}
problems.extend(validate::check(&settings));
if problems.is_empty() {
Ok(settings)
} else {
Err(ConfigError::Invalid(problems))
}
}
#[must_use]
pub fn env_keys(env: &EnvSource) -> BTreeSet<String> {
let prefix = format!("{ENV_PREFIX}_").to_lowercase();
env_pairs(env.clone())
.into_iter()
.filter(|(key, _)| !is_not_config(key))
.map(|(key, _)| key.to_lowercase())
.filter_map(|key| {
let rest = key.strip_prefix(&prefix)?;
(!rest.is_empty()).then(|| rest.replace("__", "."))
})
.collect()
}
#[must_use]
pub fn ssh_client_env(env: &EnvSource) -> Option<String> {
env_pairs(env.clone())
.into_iter()
.find(|(key, _)| key.eq_ignore_ascii_case(super::SSH_CLIENT_ENV))
.map(|(_, value)| value)
}
pub fn process_client(configured: super::SshClient) -> Result<super::SshClient, String> {
super::effective_client(configured, ssh_client_env(&EnvSource::Process).as_deref())
}
fn is_not_config(key: &str) -> bool {
let key = key.to_lowercase();
NOT_CONFIG_VARIABLES.iter().any(|name| {
let name = name.to_lowercase();
if name.ends_with('_') {
key.starts_with(&name)
} else {
key == name
}
})
}
fn env_pairs(env: EnvSource) -> Vec<(String, String)> {
match env {
EnvSource::Process => std::env::vars_os()
.filter_map(|(key, value)| Some((key.into_string().ok()?, value.into_string().ok()?)))
.collect(),
EnvSource::Vars(pairs) => pairs,
}
}
fn coerce_env_scalars(extra: &mut BTreeMap<String, Value>, file: &BTreeMap<String, Value>) {
for (key, value) in extra.iter_mut() {
coerce(value, file.get(key));
}
}
fn coerce(value: &mut Value, file: Option<&Value>) {
if let Value::Object(map) = value {
for (key, inner) in map.iter_mut() {
coerce(inner, file.and_then(|file| file.get(key)));
}
return;
}
let parsed = match &*value {
Value::String(text) if file.and_then(Value::as_str) != Some(text.as_str()) => scalar(text),
_ => None,
};
if let Some(parsed) = parsed {
*value = parsed;
}
}
fn scalar(text: &str) -> Option<Value> {
match text {
"true" => return Some(Value::Bool(true)),
"false" => return Some(Value::Bool(false)),
_ => {},
}
if let Some(integer) = canonical_integer(text) {
return Some(integer);
}
if looks_like_integer(text) {
return None;
}
text.parse::<f64>()
.ok()
.and_then(serde_json::Number::from_f64)
.map(Value::Number)
}
fn canonical_integer(text: &str) -> Option<Value> {
if let Ok(integer) = text.parse::<i64>() {
return (integer.to_string() == text).then_some(Value::from(integer));
}
if let Ok(integer) = text.parse::<u64>() {
return (integer.to_string() == text).then_some(Value::from(integer));
}
None
}
fn looks_like_integer(text: &str) -> bool {
let digits = text.strip_prefix(['+', '-']).unwrap_or(text);
!digits.is_empty() && digits.bytes().all(|byte| byte.is_ascii_digit())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn scalars_parse_only_plain_values() {
assert_eq!(scalar("true"), Some(Value::Bool(true)));
assert_eq!(scalar("10"), Some(Value::from(10)));
assert_eq!(scalar("0.05"), Some(Value::from(0.05)));
assert_eq!(scalar("nan"), None);
assert_eq!(scalar("inf"), None);
assert_eq!(scalar("qwen3"), None);
assert_eq!(scalar("True"), None);
}
#[test]
fn non_canonical_integers_stay_strings() {
assert_eq!(scalar("0123"), None, "leading zero");
assert_eq!(scalar("+5"), None, "leading plus");
assert_eq!(scalar("007"), None, "leading zeros");
}
#[test]
fn canonical_integers_and_floats_still_parse() {
assert_eq!(scalar("0"), Some(Value::from(0)));
assert_eq!(scalar("-3"), Some(Value::from(-3)));
assert_eq!(scalar("1e-7"), Some(Value::from(1e-7)));
}
#[test]
fn i64_max_still_parses_as_an_integer() {
assert_eq!(scalar("9223372036854775807"), Some(Value::from(i64::MAX)));
}
#[test]
fn i64_max_plus_one_becomes_a_u64_not_a_float() {
assert_eq!(
scalar("9223372036854775808"),
Some(Value::from(9_223_372_036_854_775_808_u64)),
"one past i64::MAX must not lose precision as an f64"
);
}
#[test]
fn u64_max_still_parses_as_an_integer() {
assert_eq!(scalar("18446744073709551615"), Some(Value::from(u64::MAX)));
}
#[test]
fn u64_max_plus_one_stays_a_string() {
assert_eq!(
scalar("18446744073709551616"),
None,
"an integer past u64::MAX must stay a string, not become an imprecise float"
);
}
#[test]
fn a_negative_integer_below_i64_min_stays_a_string() {
assert_eq!(
scalar("-9223372036854775809"),
None,
"a negative integer past i64::MIN must stay a string, not become an imprecise float"
);
}
fn vars(pairs: &[(&str, &str)]) -> EnvSource {
EnvSource::Vars(
pairs
.iter()
.map(|(k, v)| ((*k).to_string(), (*v).to_string()))
.collect(),
)
}
#[test]
fn env_keys_maps_double_underscores_to_dots_and_lower_cases() {
let env = vars(&[("OVERBRAINER_PROVIDERS__OPENROUTER__API_KEY", "sk-something")]);
assert_eq!(
env_keys(&env),
BTreeSet::from(["providers.openrouter.api_key".to_string()])
);
}
#[test]
fn env_keys_maps_log_to_a_bare_key() {
let env = vars(&[("OVERBRAINER_LOG", "debug")]);
assert_eq!(env_keys(&env), BTreeSet::from(["log".to_string()]));
}
#[test]
fn the_ssh_client_variable_is_read_apart_from_the_configuration() {
let env = vars(&[("OVERBRAINER_SSH_CLIENT", "builtin")]);
assert_eq!(env_keys(&env), BTreeSet::new());
assert_eq!(ssh_client_env(&env), Some("builtin".to_string()));
assert_eq!(ssh_client_env(&vars(&[])), None);
}
#[test]
fn the_ssh_client_variable_is_read_in_any_case() {
let env = vars(&[("overbrainer_ssh_client", "builtin")]);
assert_eq!(ssh_client_env(&env), Some("builtin".to_string()));
}
#[test]
fn env_keys_skips_tui_variables_and_unrelated_prefixes() {
let env = vars(&[
("OVERBRAINER_TUI_COLOR", "256"),
("OVERBRAINER_TUI_MOTION", "off"),
("PATH", "/usr/bin"),
("OVERBRAINERX", "not us"),
]);
assert!(env_keys(&env).is_empty());
}
#[test]
fn env_keys_skips_the_variables_read_outside_the_configuration() {
let env = vars(&[
("OVERBRAINER_NO_UPDATE_CHECK", "1"),
("overbrainer_no_update_check", "1"),
("OVERBRAINER_TEST_SSH_HOST", "overbrainer-test"),
]);
assert!(env_keys(&env).is_empty());
}
#[test]
fn env_keys_skips_the_bare_prefix_with_nothing_after_it() {
let env = vars(&[("OVERBRAINER_", "")]);
assert!(env_keys(&env).is_empty());
}
#[test]
fn env_keys_matches_a_lower_case_variable_name() {
let env = vars(&[("overbrainer_log", "debug")]);
assert_eq!(env_keys(&env), BTreeSet::from(["log".to_string()]));
}
#[test]
fn env_keys_matches_a_mixed_case_variable_name() {
let env = vars(&[("Overbrainer_Providers__Openrouter__Api_Key", "sk-something")]);
assert_eq!(
env_keys(&env),
BTreeSet::from(["providers.openrouter.api_key".to_string()])
);
}
#[test]
fn env_keys_skips_a_lower_case_tui_variable_case_insensitively() {
let env = vars(&[("overbrainer_tui_color", "256")]);
assert!(env_keys(&env).is_empty());
}
}