diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..3c063d1 --- /dev/null +++ b/src/config.rs @@ -0,0 +1,240 @@ +use crate::budget::CreditTier; +use crate::persona::{Persona, PersonaRegistry}; +use crate::pipeline::PipelineConfig; +use serde::{Deserialize, Serialize}; +use std::fs; +use std::path::{Path, PathBuf}; + +const KEYRING_SERVICE: &str = "brokecrew"; +const KEYRING_USER: &str = "openrouter_api_key"; +const DEFAULT_PERSONAS_TOML: &str = include_str!("../config/personas.example.toml"); + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AppConfigFile { + #[serde(default = "default_persona_str")] + pub default_persona: String, + #[serde(default = "default_router_model")] + pub router_model: String, + #[serde(default = "default_complexity_threshold")] + pub complexity_threshold: u8, + #[serde(default)] + pub allow_stealth: bool, + #[serde(default = "default_credit_tier")] + pub credit_tier: CreditTier, + #[serde(default = "default_base_url")] + pub openrouter_base_url: String, + pub api_key: Option, +} + +fn default_persona_str() -> String { + "general".to_string() +} +fn default_router_model() -> String { + "meta-llama/llama-3.2-3b-instruct:free".to_string() +} +fn default_complexity_threshold() -> u8 { + 7 +} +fn default_credit_tier() -> CreditTier { + CreditTier::LowBalance +} +fn default_base_url() -> String { + "https://openrouter.ai/api/v1".to_string() +} + +impl Default for AppConfigFile { + fn default() -> Self { + Self { + default_persona: default_persona_str(), + router_model: default_router_model(), + complexity_threshold: default_complexity_threshold(), + allow_stealth: false, + credit_tier: default_credit_tier(), + openrouter_base_url: default_base_url(), + api_key: None, + } + } +} + +#[derive(Debug, Deserialize, Default)] +pub struct PersonasConfigFile { + #[serde(rename = "persona", default)] + pub personas: Vec, + #[serde(rename = "pipeline", default)] + pub pipelines: Vec, +} + +pub struct AppConfig { + pub default_persona: String, + pub router_model: String, + pub complexity_threshold: u8, + pub allow_stealth: bool, + pub credit_tier: CreditTier, + pub openrouter_base_url: String, + pub api_key: Option, + pub persona_registry: PersonaRegistry, + pub pipelines: Vec, + pub config_dir: PathBuf, +} + +impl AppConfig { + pub fn default_config_dir() -> PathBuf { + dirs::config_dir() + .map(|p| p.join("brokecrew")) + .unwrap_or_else(|| PathBuf::from(".brokecrew")) + } + + pub fn load_from_paths( + config_path: Option<&Path>, + personas_path: Option<&Path>, + cli_api_key: Option, + ) -> Result { + let config_dir = Self::default_config_dir(); + + // 1. Load config.toml + let config_file_path = config_path + .map(PathBuf::from) + .unwrap_or_else(|| config_dir.join("config.toml")); + + let app_file = if config_file_path.exists() { + let data = fs::read_to_string(&config_file_path) + .map_err(|e| format!("Failed to read {}: {e}", config_file_path.display()))?; + toml::from_str::(&data) + .map_err(|e| format!("Failed to parse {}: {e}", config_file_path.display()))? + } else { + AppConfigFile::default() + }; + + // 2. Load personas.toml + let personas_file_path = personas_path + .map(PathBuf::from) + .unwrap_or_else(|| config_dir.join("personas.toml")); + + let personas_config = if personas_file_path.exists() { + let data = fs::read_to_string(&personas_file_path) + .map_err(|e| format!("Failed to read {}: {e}", personas_file_path.display()))?; + toml::from_str::(&data) + .map_err(|e| format!("Failed to parse {}: {e}", personas_file_path.display()))? + } else { + // Fallback to built-in personas.example.toml content + toml::from_str::(DEFAULT_PERSONAS_TOML) + .map_err(|e| format!("Failed to parse bundled personas: {e}"))? + }; + + // 3. Resolve API Key: CLI -> Env -> Keyring -> Config file + let api_key = Self::resolve_api_key(cli_api_key, app_file.api_key, true); + + Ok(Self { + default_persona: app_file.default_persona, + router_model: app_file.router_model, + complexity_threshold: app_file.complexity_threshold, + allow_stealth: app_file.allow_stealth, + credit_tier: app_file.credit_tier, + openrouter_base_url: app_file.openrouter_base_url, + api_key, + persona_registry: PersonaRegistry::new(personas_config.personas), + pipelines: personas_config.pipelines, + config_dir, + }) + } + + /// Resolves the API key using precedence: + /// 1. CLI argument + /// 2. `OPENROUTER_API_KEY` environment variable + /// 3. OS Keyring + /// 4. Config file + pub fn resolve_api_key( + cli_key: Option, + config_file_key: Option, + use_keyring: bool, + ) -> Option { + if let Some(key) = cli_key.filter(|k| !k.trim().is_empty()) { + return Some(key.trim().to_string()); + } + + if let Ok(env_key) = std::env::var("OPENROUTER_API_KEY") { + if !env_key.trim().is_empty() { + return Some(env_key.trim().to_string()); + } + } + + if use_keyring { + if let Ok(entry) = keyring::Entry::new(KEYRING_SERVICE, KEYRING_USER) { + if let Ok(secret) = entry.get_password() { + if !secret.trim().is_empty() { + return Some(secret.trim().to_string()); + } + } + } + } + + if let Some(key) = config_file_key.filter(|k| !k.trim().is_empty()) { + return Some(key.trim().to_string()); + } + + None + } + + /// Saves API key securely to the OS credential store + pub fn save_api_key_to_keyring(api_key: &str) -> Result<(), String> { + let entry = keyring::Entry::new(KEYRING_SERVICE, KEYRING_USER) + .map_err(|e| format!("Failed to access OS keyring: {e}"))?; + entry + .set_password(api_key.trim()) + .map_err(|e| format!("Failed to store secret in OS keyring: {e}")) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_resolve_api_key_priority() { + // 1. CLI flag has highest priority + let key = AppConfig::resolve_api_key( + Some("cli-key-123".into()), + Some("file-key-456".into()), + false, + ); + assert_eq!(key, Some("cli-key-123".into())); + + // 2. Fallback to config file when no env or cli + let key = AppConfig::resolve_api_key(None, Some("file-key-456".into()), false); + assert_eq!(key, Some("file-key-456".into())); + } + + #[test] + fn test_personas_config_file_parse() { + let toml_str = r#" + [[persona]] + name = "tester" + system_prompt = "You test code" + trigger_words = ["test", "verify"] + + [[persona.candidates]] + id = "model-a:free" + tier = "free" + is_stealth = false + + [[pipeline]] + name = "test-pipe" + steps = [ + { persona = "tester" } + ] + "#; + + let parsed: PersonasConfigFile = toml::from_str(toml_str).unwrap(); + assert_eq!(parsed.personas.len(), 1); + assert_eq!(parsed.personas[0].name, "tester"); + assert_eq!(parsed.pipelines.len(), 1); + assert_eq!(parsed.pipelines[0].name, "test-pipe"); + } + + #[test] + fn test_bundled_personas_valid() { + let parsed: PersonasConfigFile = toml::from_str(DEFAULT_PERSONAS_TOML).unwrap(); + assert!(parsed.personas.len() >= 3); + assert!(parsed.pipelines.len() >= 1); + } +} diff --git a/src/lib.rs b/src/lib.rs index 87b277c..cd2a1b7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,6 +2,7 @@ pub mod budget; pub mod client; +pub mod config; pub mod persona; pub mod pipeline; pub mod registry; @@ -9,6 +10,7 @@ pub mod router; pub use budget::{BudgetExceeded, BudgetStatus, BudgetTracker, CreditTier}; pub use client::{ClientError, ModelClient, OpenRouterClient}; +pub use config::{AppConfig, AppConfigFile, PersonasConfigFile}; pub use persona::{DispatchResult, ModelAttempt, ModelCandidate, Persona, PersonaRegistry, Tier}; pub use pipeline::{PipelineConfig, PipelineEvent, PipelineExecutor, PipelineStep}; pub use registry::{ModelInfo, ModelPricing, ModelRegistry};