diff --git a/src/budget.rs b/src/budget.rs index 3083c11..5addc52 100644 --- a/src/budget.rs +++ b/src/budget.rs @@ -163,8 +163,7 @@ impl BudgetTracker { last_day_utc: self.last_day_utc.clone(), recent_calls: self.recent_calls.clone(), }; - let serialized = serde_json::to_string_pretty(&state) - .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?; + let serialized = serde_json::to_string_pretty(&state).map_err(std::io::Error::other)?; // Atomic write via tempfile in the same parent directory let parent_dir = path.parent().unwrap_or_else(|| Path::new(".")); diff --git a/src/cli.rs b/src/cli.rs index e71c11e..a59393e 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -15,10 +15,19 @@ pub struct Cli { #[arg(long, global = true, help = "Path to personas.toml file")] pub personas_file: Option, - #[arg(long, global = true, env = "OPENROUTER_API_KEY", help = "OpenRouter API Key (overrides keyring & env)")] + #[arg( + long, + global = true, + env = "OPENROUTER_API_KEY", + help = "OpenRouter API Key (overrides keyring & env)" + )] pub api_key: Option, - #[arg(long, global = true, help = "Output structured logs in JSON format to stderr")] + #[arg( + long, + global = true, + help = "Output structured logs in JSON format to stderr" + )] pub json_logs: bool, #[arg(long, short, global = true, help = "Enable verbose debug logs")] diff --git a/src/client.rs b/src/client.rs index ef504ae..7f50078 100644 --- a/src/client.rs +++ b/src/client.rs @@ -136,7 +136,11 @@ impl OpenRouterClient { None } - fn check_status(status: reqwest::StatusCode, body: &str, headers: &reqwest::header::HeaderMap) -> Result<(), ClientError> { + fn check_status( + status: reqwest::StatusCode, + body: &str, + headers: &reqwest::header::HeaderMap, + ) -> Result<(), ClientError> { match status.as_u16() { 200..=299 => Ok(()), 401 => Err(ClientError::Unauthorized), @@ -149,7 +153,9 @@ impl OpenRouterClient { .unwrap_or(5); Err(ClientError::RateLimited(retry_after)) } - 404 => Err(ClientError::ModelUnavailable(format!("Model not found (404): {body}"))), + 404 => Err(ClientError::ModelUnavailable(format!( + "Model not found (404): {body}" + ))), s if s >= 500 => Err(ClientError::ServerError { status: s, message: body.to_string(), @@ -223,7 +229,8 @@ impl OpenRouterClient { #[async_trait] impl ModelClient for OpenRouterClient { async fn complete(&self, model: &str, system: &str, user: &str) -> Result { - self.complete_with_params(model, system, user, None, None).await + self.complete_with_params(model, system, user, None, None) + .await } async fn complete_stream( @@ -276,7 +283,7 @@ impl ModelClient for OpenRouterClient { let mut buffer = String::new(); while let Some(item) = stream.next().await { - let bytes = item.map_err(|e| ClientError::Network(e))?; + let bytes = item.map_err(ClientError::Network)?; let chunk_str = String::from_utf8_lossy(&bytes); buffer.push_str(&chunk_str); @@ -331,27 +338,47 @@ mod tests { #[test] fn test_check_status_mappings() { let mut headers = reqwest::header::HeaderMap::new(); - headers.insert("retry-after", reqwest::header::HeaderValue::from_static("12")); + headers.insert( + "retry-after", + reqwest::header::HeaderValue::from_static("12"), + ); - let err_401 = OpenRouterClient::check_status(reqwest::StatusCode::UNAUTHORIZED, "bad key", &headers).unwrap_err(); + let err_401 = + OpenRouterClient::check_status(reqwest::StatusCode::UNAUTHORIZED, "bad key", &headers) + .unwrap_err(); match err_401 { ClientError::Unauthorized => (), _ => panic!("Expected Unauthorized, got {:?}", err_401), } - let err_402 = OpenRouterClient::check_status(reqwest::StatusCode::PAYMENT_REQUIRED, "low balance", &headers).unwrap_err(); + let err_402 = OpenRouterClient::check_status( + reqwest::StatusCode::PAYMENT_REQUIRED, + "low balance", + &headers, + ) + .unwrap_err(); match err_402 { ClientError::PaymentRequired => (), _ => panic!("Expected PaymentRequired, got {:?}", err_402), } - let err_429 = OpenRouterClient::check_status(reqwest::StatusCode::TOO_MANY_REQUESTS, "rate limited", &headers).unwrap_err(); + let err_429 = OpenRouterClient::check_status( + reqwest::StatusCode::TOO_MANY_REQUESTS, + "rate limited", + &headers, + ) + .unwrap_err(); match err_429 { ClientError::RateLimited(secs) => assert_eq!(secs, 12), _ => panic!("Expected RateLimited(12), got {:?}", err_429), } - let err_503 = OpenRouterClient::check_status(reqwest::StatusCode::SERVICE_UNAVAILABLE, "outage", &headers).unwrap_err(); + let err_503 = OpenRouterClient::check_status( + reqwest::StatusCode::SERVICE_UNAVAILABLE, + "outage", + &headers, + ) + .unwrap_err(); match err_503 { ClientError::ServerError { status, .. } => assert_eq!(status, 503), _ => panic!("Expected ServerError 503, got {:?}", err_503), diff --git a/src/config.rs b/src/config.rs index 3c063d1..b16fcdc 100644 --- a/src/config.rs +++ b/src/config.rs @@ -235,6 +235,6 @@ mod tests { 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); + assert!(!parsed.pipelines.is_empty()); } } diff --git a/src/main.rs b/src/main.rs index 47941bb..2469fd7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -103,7 +103,12 @@ async fn main() -> Result<(), Box> { AppConfig::save_api_key_to_keyring(&api_key) .map_err(|e| format!("Failed to save API key to OS credential store: {e}"))?; - println!("{}", "✔ OpenRouter API key saved securely in OS credential store.".bright_green().bold()); + println!( + "{}", + "✔ OpenRouter API key saved securely in OS credential store." + .bright_green() + .bold() + ); return Ok(()); } }, @@ -115,7 +120,10 @@ async fn main() -> Result<(), Box> { let badge = get_persona_color_tag(&persona.name); println!("\n {} - {}", badge, persona.system_prompt.dimmed()); if !persona.trigger_words.is_empty() { - println!(" Triggers: {}", persona.trigger_words.join(", ").italic()); + println!( + " Triggers: {}", + persona.trigger_words.join(", ").italic() + ); } println!(" Candidates:"); for (i, c) in persona.candidates.iter().enumerate() { @@ -128,7 +136,13 @@ async fn main() -> Result<(), Box> { } else { String::new() }; - println!(" {}. {} ({}){}", i + 1, c.id.bold(), tier_tag, stealth_tag); + println!( + " {}. {} ({}){}", + i + 1, + c.id.bold(), + tier_tag, + stealth_tag + ); } } @@ -146,7 +160,11 @@ async fn main() -> Result<(), Box> { } }) .collect(); - println!(" • {} ➔ {}", pipe.name.bright_cyan().bold(), steps_str.join(" ➔ ")); + println!( + " • {} ➔ {}", + pipe.name.bright_cyan().bold(), + steps_str.join(" ➔ ") + ); } } println!(); @@ -171,20 +189,39 @@ async fn main() -> Result<(), Box> { }; if model_count > 0 { - println!(" Live catalog: {} available models cached", model_count.to_string().bright_green().bold()); + println!( + " Live catalog: {} available models cached", + model_count.to_string().bright_green().bold() + ); } match budget_tracker.check_budget() { Ok(status) => { let pct = (status.daily_requests * 100) / status.daily_limit; let daily_color = if pct >= 80 { - format!("{}/{} ({}%)", status.daily_requests, status.daily_limit, pct).yellow().bold().to_string() + format!( + "{}/{} ({}%)", + status.daily_requests, status.daily_limit, pct + ) + .yellow() + .bold() + .to_string() } else { - format!("{}/{} ({}%)", status.daily_requests, status.daily_limit, pct).green().bold().to_string() + format!( + "{}/{} ({}%)", + status.daily_requests, status.daily_limit, pct + ) + .green() + .bold() + .to_string() }; println!(" Account tier: {:?}", config.credit_tier); - println!(" Rolling 1-min requests: {}/{}", status.minute_requests.to_string().bold(), status.minute_limit); + println!( + " Rolling 1-min requests: {}/{}", + status.minute_requests.to_string().bold(), + status.minute_limit + ); println!(" Daily requests today: {}", daily_color); if status.approaching_daily_limit { println!(" {}", "⚠ Notice: Approaching daily free-tier limit. Non-urgent calls may be deferred.".yellow().bold()); @@ -204,7 +241,13 @@ async fn main() -> Result<(), Box> { None => { // Default to chat if no subcommand passed - run_chat(ChatArgs::default(), &config, &mut budget_tracker, model_registry).await?; + run_chat( + ChatArgs::default(), + &config, + &mut budget_tracker, + model_registry, + ) + .await?; } } @@ -261,8 +304,14 @@ async fn run_chat( } // Otherwise, run interactive REPL loop - println!("\n{}", "BrokeCrew Interactive Team Chat".bright_cyan().bold()); - println!("{}", "Type your request and press Enter. Type 'exit' or Ctrl+C to quit.\n".dimmed()); + println!( + "\n{}", + "BrokeCrew Interactive Team Chat".bright_cyan().bold() + ); + println!( + "{}", + "Type your request and press Enter. Type 'exit' or Ctrl+C to quit.\n".dimmed() + ); let stdin = io::stdin(); let mut reader = stdin.lock(); @@ -375,14 +424,20 @@ async fn execute_prompt( (dec.persona, Some(info)) } Err(e) => { - eprintln!("Router warning: {}, falling back to {}", e, config.default_persona); + eprintln!( + "Router warning: {}, falling back to {}", + e, config.default_persona + ); (config.default_persona.clone(), None) } } }; if let Some(info) = route_info { - eprintln!("{}", format!("⚡ [BrokeCrew: {info} ➔ {}]", target_persona).dimmed()); + eprintln!( + "{}", + format!("⚡ [BrokeCrew: {info} ➔ {}]", target_persona).dimmed() + ); } let (chunk_tx, mut chunk_rx) = mpsc::channel::(100); @@ -435,6 +490,7 @@ async fn execute_prompt( Ok(()) } +#[allow(clippy::too_many_arguments)] async fn run_pipeline( pipeline: &PipelineConfig, config: &AppConfig, diff --git a/src/persona.rs b/src/persona.rs index 9390016..ea3deb2 100644 --- a/src/persona.rs +++ b/src/persona.rs @@ -49,7 +49,15 @@ impl PersonaRegistry { pub fn new(personas: Vec) -> Self { Self { personas } } +} +impl From> for PersonaRegistry { + fn from(personas: Vec) -> Self { + Self::new(personas) + } +} + +impl PersonaRegistry { pub fn get(&self, name: &str) -> Option<&Persona> { self.personas.iter().find(|p| p.name == name) } @@ -152,6 +160,8 @@ mod tests { assert_eq!(paid_only[0].id, "gemini-paid"); // Non-existent persona returns error - assert!(registry.get_eligible_candidates("unknown", false, false).is_err()); + assert!(registry + .get_eligible_candidates("unknown", false, false) + .is_err()); } } diff --git a/src/pipeline.rs b/src/pipeline.rs index 2cca241..e346701 100644 --- a/src/pipeline.rs +++ b/src/pipeline.rs @@ -45,6 +45,7 @@ pub struct PipelineExecutor; impl PipelineExecutor { /// Executes a single persona query with candidate fallback walk + #[allow(clippy::too_many_arguments)] pub async fn dispatch_persona( persona_name: &str, persona_registry: &PersonaRegistry, @@ -55,9 +56,9 @@ impl PipelineExecutor { require_paid: bool, chunk_sender: Option>, ) -> Result { - let persona = persona_registry - .get(persona_name) - .ok_or_else(|| ClientError::ModelUnavailable(format!("Persona '{persona_name}' not found")))?; + let persona = persona_registry.get(persona_name).ok_or_else(|| { + ClientError::ModelUnavailable(format!("Persona '{persona_name}' not found")) + })?; // Filter candidates against live registry and options let eligible = model_registry @@ -68,7 +69,11 @@ impl PipelineExecutor { eligible } else { // If live registry filtering yielded nothing (e.g. offline/strict), fall back to raw persona candidates - persona.candidates.iter().filter(|c| allow_stealth || !c.is_stealth).collect() + persona + .candidates + .iter() + .filter(|c| allow_stealth || !c.is_stealth) + .collect() }; if candidates_to_try.is_empty() { @@ -166,6 +171,7 @@ impl PipelineExecutor { } /// Executes an entire pipeline sequentially, passing outputs between steps + #[allow(clippy::too_many_arguments)] pub async fn execute( pipeline: &PipelineConfig, persona_registry: &PersonaRegistry, @@ -181,9 +187,9 @@ impl PipelineExecutor { let total_steps = pipeline.steps.len(); for (step_idx, step) in pipeline.steps.iter().enumerate() { - let persona = persona_registry - .get(&step.persona) - .ok_or_else(|| ClientError::ModelUnavailable(format!("Persona '{}' not found", step.persona)))?; + let persona = persona_registry.get(&step.persona).ok_or_else(|| { + ClientError::ModelUnavailable(format!("Persona '{}' not found", step.persona)) + })?; // Prepare prompt for this step let step_prompt = if step.pass_previous_output && !previous_output.is_empty() { @@ -286,7 +292,12 @@ mod tests { #[async_trait] impl ModelClient for MockPipelineClient { - async fn complete(&self, model: &str, _system: &str, user: &str) -> Result { + async fn complete( + &self, + model: &str, + _system: &str, + user: &str, + ) -> Result { self.calls.fetch_add(1, Ordering::SeqCst); Ok(format!("Output from {} for: {}", model, user)) } diff --git a/src/registry.rs b/src/registry.rs index 7745470..d2ca24a 100644 --- a/src/registry.rs +++ b/src/registry.rs @@ -98,7 +98,10 @@ impl ModelRegistry { .map_err(|e| format!("Network error fetching models catalog: {e}"))?; if !res.status().is_success() { - return Err(format!("OpenRouter models endpoint returned HTTP {}", res.status())); + return Err(format!( + "OpenRouter models endpoint returned HTTP {}", + res.status() + )); } let body = res @@ -184,10 +187,17 @@ mod tests { ] }"#; - let count = registry.update_from_json(json_data).await.expect("update failed"); + let count = registry + .update_from_json(json_data) + .await + .expect("update failed"); assert_eq!(count, 2); - assert!(registry.is_model_available("qwen/qwen-2.5-coder-32b-instruct:free").await); + assert!( + registry + .is_model_available("qwen/qwen-2.5-coder-32b-instruct:free") + .await + ); assert!(!registry.is_model_available("unlisted/model").await); let candidates = vec![ diff --git a/src/router/keyword.rs b/src/router/keyword.rs index 9548c46..c03ffc6 100644 --- a/src/router/keyword.rs +++ b/src/router/keyword.rs @@ -70,13 +70,24 @@ mod tests { Persona { name: "coder".into(), system_prompt: "coder prompt".into(), - trigger_words: vec!["build".into(), "code".into(), "script".into(), "website".into(), "api".into()], + trigger_words: vec![ + "build".into(), + "code".into(), + "script".into(), + "website".into(), + "api".into(), + ], candidates: vec![], }, Persona { name: "reviewer".into(), system_prompt: "reviewer prompt".into(), - trigger_words: vec!["review".into(), "audit".into(), "critique".into(), "lint".into()], + trigger_words: vec![ + "review".into(), + "audit".into(), + "critique".into(), + "lint".into(), + ], candidates: vec![], }, Persona { @@ -97,7 +108,10 @@ mod tests { ("write a python script to parse logs", Some("coder")), ("please review this pull request for bugs", Some("reviewer")), ("audit our authentication system", Some("reviewer")), - ("design a high scale distributed cache architecture", Some("architect")), + ( + "design a high scale distributed cache architecture", + Some("architect"), + ), ("hello, what is the weather today?", None), ("how are you doing?", None), ]; diff --git a/src/router/llm.rs b/src/router/llm.rs index 585b81b..96ed7c9 100644 --- a/src/router/llm.rs +++ b/src/router/llm.rs @@ -15,9 +15,15 @@ impl LlmRouter { known_personas.join(", ") ); - let user_prompt = format!("Classify this request to the most appropriate persona:\n\"{}\"", prompt); + let user_prompt = format!( + "Classify this request to the most appropriate persona:\n\"{}\"", + prompt + ); - let raw_output = match client.complete(router_model, &system_prompt, &user_prompt).await { + let raw_output = match client + .complete(router_model, &system_prompt, &user_prompt) + .await + { Ok(out) => out, Err(e) => { tracing::warn!("Router LLM call failed, falling back to default persona: {e}"); @@ -26,9 +32,15 @@ impl LlmRouter { }; let cleaned = Self::sanitize_output(&raw_output); - if known_personas.iter().any(|p| p.eq_ignore_ascii_case(&cleaned)) { + if known_personas + .iter() + .any(|p| p.eq_ignore_ascii_case(&cleaned)) + { // Find matched casing - if let Some(matched) = known_personas.iter().find(|p| p.eq_ignore_ascii_case(&cleaned)) { + if let Some(matched) = known_personas + .iter() + .find(|p| p.eq_ignore_ascii_case(&cleaned)) + { return Ok(matched.clone()); } } @@ -94,7 +106,12 @@ mod tests { #[async_trait] impl ModelClient for MockRouterClient { - async fn complete(&self, _model: &str, _system: &str, _user: &str) -> Result { + async fn complete( + &self, + _model: &str, + _system: &str, + _user: &str, + ) -> Result { Ok(self.reply.clone()) } @@ -117,27 +134,45 @@ mod tests { let client_valid = MockRouterClient { reply: "coder".into(), }; - let res = LlmRouter::classify(&client_valid, "router-model", &known, "general", "build something") - .await - .unwrap(); + let res = LlmRouter::classify( + &client_valid, + "router-model", + &known, + "general", + "build something", + ) + .await + .unwrap(); assert_eq!(res, "coder"); // 2. Extra whitespace & quotes let client_messy = MockRouterClient { reply: " \"REVIEWER\" \n".into(), }; - let res = LlmRouter::classify(&client_messy, "router-model", &known, "general", "audit this") - .await - .unwrap(); + let res = LlmRouter::classify( + &client_messy, + "router-model", + &known, + "general", + "audit this", + ) + .await + .unwrap(); assert_eq!(res, "reviewer"); // 3. Hallucinated / unrecognized persona defaults to default_persona let client_hallucinated = MockRouterClient { reply: "I believe a poet would be best for this query".into(), }; - let res = LlmRouter::classify(&client_hallucinated, "router-model", &known, "general", "write poem") - .await - .unwrap(); + let res = LlmRouter::classify( + &client_hallucinated, + "router-model", + &known, + "general", + "write poem", + ) + .await + .unwrap(); assert_eq!(res, "general"); } diff --git a/src/router/mod.rs b/src/router/mod.rs index df8cf5a..324d208 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -87,7 +87,12 @@ mod tests { #[async_trait] impl ModelClient for DummyClient { - async fn complete(&self, _model: &str, _system: &str, _user: &str) -> Result { + async fn complete( + &self, + _model: &str, + _system: &str, + _user: &str, + ) -> Result { Ok("general".into()) } async fn complete_stream( @@ -118,17 +123,27 @@ mod tests { }, ]; let registry = PersonaRegistry::new(personas); - let router = Router::new("general".into(), "meta-llama/llama-3.2-3b-instruct:free".into(), 7); + let router = Router::new( + "general".into(), + "meta-llama/llama-3.2-3b-instruct:free".into(), + 7, + ); let client = DummyClient; // Matches Stage A (build) without escalation - let dec1 = router.route(®istry, &client, "build a script", false).await.unwrap(); + let dec1 = router + .route(®istry, &client, "build a script", false) + .await + .unwrap(); assert_eq!(dec1.persona, "coder"); assert_eq!(dec1.stage_used, RouteStage::StageAKeyword); assert!(!dec1.escalate_to_paid); // Matches Stage A with forced escalation - let dec2 = router.route(®istry, &client, "build a script", true).await.unwrap(); + let dec2 = router + .route(®istry, &client, "build a script", true) + .await + .unwrap(); assert_eq!(dec2.persona, "coder"); assert!(dec2.escalate_to_paid); } diff --git a/tests/fallback_tests.rs b/tests/fallback_tests.rs new file mode 100644 index 0000000..324ae87 --- /dev/null +++ b/tests/fallback_tests.rs @@ -0,0 +1,174 @@ +use async_trait::async_trait; +use brokecrew::client::{ClientError, ModelClient}; +use brokecrew::persona::{ModelCandidate, Persona, PersonaRegistry, Tier}; +use brokecrew::pipeline::PipelineExecutor; +use brokecrew::registry::ModelRegistry; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use tokio::sync::mpsc; + +struct FaultInjectingClient { + candidate1_error: Option, + calls: Arc, +} + +#[async_trait] +impl ModelClient for FaultInjectingClient { + async fn complete(&self, model: &str, system: &str, user: &str) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + if model == "candidate-1" { + if let Some(ref err) = self.candidate1_error { + return match err { + ClientError::RateLimited(s) => Err(ClientError::RateLimited(*s)), + ClientError::ServerError { status, message } => Err(ClientError::ServerError { + status: *status, + message: message.clone(), + }), + ClientError::Unauthorized => Err(ClientError::Unauthorized), + ClientError::PaymentRequired => Err(ClientError::PaymentRequired), + _ => Err(ClientError::ModelUnavailable("simulated".into())), + }; + } + } + + Ok(format!( + "Success from {model} for prompt '{user}' (system: {system})" + )) + } + + async fn complete_stream( + &self, + model: &str, + system: &str, + user: &str, + chunk_sender: mpsc::Sender, + ) -> Result { + let reply = self.complete(model, system, user).await?; + let _ = chunk_sender.send(reply.clone()).await; + Ok(reply) + } +} + +fn create_resilience_persona() -> Persona { + Persona { + name: "resilient-coder".into(), + system_prompt: "You are coder".into(), + trigger_words: vec!["code".into()], + candidates: vec![ + ModelCandidate { + id: "candidate-1".into(), + tier: Tier::Free, + is_stealth: false, + }, + ModelCandidate { + id: "candidate-2".into(), + tier: Tier::Free, + is_stealth: false, + }, + ], + } +} + +#[tokio::test] +async fn test_candidate_cascade_on_503_service_unavailable() { + let persona = create_resilience_persona(); + let persona_reg = PersonaRegistry::new(vec![persona]); + let model_reg = ModelRegistry::new(); + + let calls = Arc::new(AtomicUsize::new(0)); + let client = FaultInjectingClient { + candidate1_error: Some(ClientError::ServerError { + status: 503, + message: "Overloaded".into(), + }), + calls: calls.clone(), + }; + + let result = PipelineExecutor::dispatch_persona( + "resilient-coder", + &persona_reg, + &model_reg, + &client, + "write rust code", + false, + false, + None, + ) + .await + .expect("fallback to candidate-2 should succeed"); + + assert_eq!(result.persona, "resilient-coder"); + assert_eq!(result.model_used, "candidate-2"); + assert!(result.fell_back, "Result must indicate fallback occurred"); + assert_eq!(result.attempts.len(), 1); + assert_eq!(result.attempts[0].model_id, "candidate-1"); + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "Both candidates must have been called" + ); +} + +#[tokio::test] +async fn test_candidate_cascade_on_429_rate_limited() { + let persona = create_resilience_persona(); + let persona_reg = PersonaRegistry::new(vec![persona]); + let model_reg = ModelRegistry::new(); + + let calls = Arc::new(AtomicUsize::new(0)); + let client = FaultInjectingClient { + candidate1_error: Some(ClientError::RateLimited(15)), + calls: calls.clone(), + }; + + let result = PipelineExecutor::dispatch_persona( + "resilient-coder", + &persona_reg, + &model_reg, + &client, + "write python code", + false, + false, + None, + ) + .await + .expect("fallback to candidate-2 should succeed"); + + assert_eq!(result.model_used, "candidate-2"); + assert!(result.fell_back); + assert_eq!(calls.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn test_fail_fast_on_401_unauthorized() { + let persona = create_resilience_persona(); + let persona_reg = PersonaRegistry::new(vec![persona]); + let model_reg = ModelRegistry::new(); + + let calls = Arc::new(AtomicUsize::new(0)); + let client = FaultInjectingClient { + candidate1_error: Some(ClientError::Unauthorized), + calls: calls.clone(), + }; + + let err = PipelineExecutor::dispatch_persona( + "resilient-coder", + &persona_reg, + &model_reg, + &client, + "test query", + false, + false, + None, + ) + .await + .unwrap_err(); + + match err { + ClientError::Unauthorized => (), + _ => panic!("Expected Unauthorized fail-fast, got {:?}", err), + } + + // Must NOT call candidate-2 when 401 Unauthorized occurs! + assert_eq!(calls.load(Ordering::SeqCst), 1); +} diff --git a/tests/router_tests.rs b/tests/router_tests.rs new file mode 100644 index 0000000..8149833 --- /dev/null +++ b/tests/router_tests.rs @@ -0,0 +1,148 @@ +use async_trait::async_trait; +use brokecrew::client::{ClientError, ModelClient}; +use brokecrew::persona::{ModelCandidate, Persona, PersonaRegistry, Tier}; +use brokecrew::router::{RouteDecision, RouteStage, Router}; +use tokio::sync::mpsc; + +struct TestRouterClient { + llm_choice: String, +} + +#[async_trait] +impl ModelClient for TestRouterClient { + async fn complete( + &self, + _model: &str, + _system: &str, + _user: &str, + ) -> Result { + Ok(self.llm_choice.clone()) + } + + async fn complete_stream( + &self, + model: &str, + system: &str, + user: &str, + _chunk_sender: mpsc::Sender, + ) -> Result { + self.complete(model, system, user).await + } +} + +fn setup_test_registry() -> PersonaRegistry { + vec![ + Persona { + name: "coder".into(), + system_prompt: "You write code".into(), + trigger_words: vec![ + "build".into(), + "code".into(), + "script".into(), + "website".into(), + "api".into(), + ], + candidates: vec![ModelCandidate { + id: "qwen-free".into(), + tier: Tier::Free, + is_stealth: false, + }], + }, + Persona { + name: "reviewer".into(), + system_prompt: "You review code".into(), + trigger_words: vec!["review".into(), "audit".into(), "critique".into()], + candidates: vec![ModelCandidate { + id: "llama-free".into(), + tier: Tier::Free, + is_stealth: false, + }], + }, + Persona { + name: "general".into(), + system_prompt: "You are general assistant".into(), + trigger_words: vec!["help".into(), "hello".into()], + candidates: vec![ModelCandidate { + id: "general-free".into(), + tier: Tier::Free, + is_stealth: false, + }], + }, + ] + .into() +} + +#[tokio::test] +async fn test_table_driven_classification() { + let registry = setup_test_registry(); + let client = TestRouterClient { + llm_choice: "reviewer".into(), + }; + let router = Router::new("general".into(), "router-model".into(), 7); + + let cases = vec![ + ( + "build me a proxy checker site", + "coder", + RouteStage::StageAKeyword, + false, + ), + ( + "please audit this auth module", + "reviewer", + RouteStage::StageAKeyword, + false, + ), + ( + "how does photosynthesis work?", + "reviewer", + RouteStage::StageBLlm, + false, + ), // LLM choice returned "reviewer" + ]; + + for (prompt, expected_persona, expected_stage, expected_escalate) in cases { + let decision: RouteDecision = router + .route(®istry, &client, prompt, false) + .await + .expect("routing should succeed"); + + assert_eq!( + decision.persona, expected_persona, + "Prompt failed: '{}'", + prompt + ); + assert_eq!( + decision.stage_used, expected_stage, + "Stage mismatch for: '{}'", + prompt + ); + assert_eq!( + decision.escalate_to_paid, expected_escalate, + "Escalate mismatch for: '{}'", + prompt + ); + } +} + +#[tokio::test] +async fn test_complexity_escalation_trigger() { + let registry = setup_test_registry(); + let client = TestRouterClient { + llm_choice: "coder".into(), + }; + let router = Router::new("general".into(), "router-model".into(), 7); + + let heavy_prompt = "Refactor all microservices across the entire project with database migration and distributed architecture:\n```rust\nfn main() {}\n```"; + let decision = router + .route(®istry, &client, heavy_prompt, false) + .await + .expect("routing should succeed"); + + assert_eq!(decision.persona, "coder"); + assert!( + decision.escalate_to_paid, + "Heavy prompt must trigger escalation" + ); + assert!(decision.complexity_score >= 7); +} diff --git a/tests/wiremock_client.rs b/tests/wiremock_client.rs new file mode 100644 index 0000000..478c9d9 --- /dev/null +++ b/tests/wiremock_client.rs @@ -0,0 +1,138 @@ +use brokecrew::client::{ClientError, ModelClient, OpenRouterClient}; +use brokecrew::registry::ModelRegistry; +use tokio::sync::mpsc; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +#[tokio::test] +async fn test_wiremock_sse_streaming() { + let mock_server = MockServer::start().await; + + let sse_body = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello \"}}]}\n\ndata: {\"choices\":[{\"delta\":{\"content\":\"from \"}}]}\n\ndata: {\"choices\":[{\"delta\":{\"content\":\"WireMock!\"}}]}\n\ndata: [DONE]\n\n"; + + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(sse_body), + ) + .mount(&mock_server) + .await; + + let client = OpenRouterClient::with_base_url("test-key".into(), mock_server.uri()); + + let (tx, mut rx) = mpsc::channel::(10); + let collect_task = tokio::spawn(async move { + let mut chunks = Vec::new(); + while let Some(c) = rx.recv().await { + chunks.push(c); + } + chunks + }); + + let full_reply = client + .complete_stream("mock-model", "system", "user", tx) + .await + .expect("streaming should succeed"); + + let chunks = collect_task.await.unwrap(); + + assert_eq!(full_reply, "Hello from WireMock!"); + assert_eq!(chunks, vec!["Hello ", "from ", "WireMock!"]); +} + +#[tokio::test] +async fn test_wiremock_401_unauthorized() { + let mock_server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(401).set_body_string("Invalid API key")) + .mount(&mock_server) + .await; + + let client = OpenRouterClient::with_base_url("bad-key".into(), mock_server.uri()); + let err = client + .complete("mock-model", "sys", "user") + .await + .unwrap_err(); + + match err { + ClientError::Unauthorized => (), + _ => panic!("Expected ClientError::Unauthorized, got {:?}", err), + } +} + +#[tokio::test] +async fn test_wiremock_402_payment_required() { + let mock_server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(402).set_body_string("Account balance depleted")) + .mount(&mock_server) + .await; + + let client = OpenRouterClient::with_base_url("key".into(), mock_server.uri()); + let err = client + .complete("mock-model", "sys", "user") + .await + .unwrap_err(); + + match err { + ClientError::PaymentRequired => (), + _ => panic!("Expected ClientError::PaymentRequired, got {:?}", err), + } +} + +#[tokio::test] +async fn test_wiremock_models_catalog_fetch() { + let mock_server = MockServer::start().await; + + let catalog_json = r#"{ + "data": [ + { + "id": "qwen/qwen-2.5-coder-32b-instruct:free", + "name": "Qwen 2.5 Coder", + "context_length": 32768, + "pricing": { "prompt": "0", "completion": "0" } + }, + { + "id": "meta-llama/llama-3.3-70b-instruct:free", + "name": "Llama 3.3 70B", + "context_length": 131072, + "pricing": { "prompt": "0", "completion": "0" } + } + ] + }"#; + + Mock::given(method("GET")) + .and(path("/models")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/json") + .set_body_string(catalog_json), + ) + .mount(&mock_server) + .await; + + let registry = ModelRegistry::new(); + let count = registry + .fetch_live(&mock_server.uri()) + .await + .expect("catalog fetch failed"); + + assert_eq!(count, 2); + assert!( + registry + .is_model_available("qwen/qwen-2.5-coder-32b-instruct:free") + .await + ); + assert!( + registry + .is_model_available("meta-llama/llama-3.3-70b-instruct:free") + .await + ); + assert!(!registry.is_model_available("non-existent-model").await); +}