From 8c7a08edfc5ce9e1125b5b30d2dae5623fa5d13d Mon Sep 17 00:00:00 2001 From: oiupoyt Date: Sat, 26 Sep 2026 14:04:26 +0530 Subject: [PATCH] fix: address code review findings on cascade failure reporting, auth fail-fast, and budget tracking --- src/client.rs | 14 +++++- src/main.rs | 7 +++ src/pipeline.rs | 35 ++++++++++++-- src/router/llm.rs | 3 ++ tests/fallback_tests.rs | 101 ++++++++++++++++++++++++++++++++++++++++ 5 files changed, 153 insertions(+), 7 deletions(-) diff --git a/src/client.rs b/src/client.rs index 7f50078..c6a87ea 100644 --- a/src/client.rs +++ b/src/client.rs @@ -194,10 +194,15 @@ impl OpenRouterClient { temperature, }; + let mut auth_val = + reqwest::header::HeaderValue::from_str(&format!("Bearer {}", self.api_key)) + .unwrap_or_else(|_| reqwest::header::HeaderValue::from_static("")); + auth_val.set_sensitive(true); + let res = self .http .post(&endpoint) - .header("Authorization", format!("Bearer {}", self.api_key)) + .header("Authorization", auth_val) .header("HTTP-Referer", "https://github.com/brokecrew/brokecrew") .header("X-Title", "BrokeCrew") .header("Content-Type", "application/json") @@ -259,10 +264,15 @@ impl ModelClient for OpenRouterClient { temperature: None, }; + let mut auth_val = + reqwest::header::HeaderValue::from_str(&format!("Bearer {}", self.api_key)) + .unwrap_or_else(|_| reqwest::header::HeaderValue::from_static("")); + auth_val.set_sensitive(true); + let res = self .http .post(&endpoint) - .header("Authorization", format!("Bearer {}", self.api_key)) + .header("Authorization", auth_val) .header("HTTP-Referer", "https://github.com/brokecrew/brokecrew") .header("X-Title", "BrokeCrew") .header("Content-Type", "application/json") diff --git a/src/main.rs b/src/main.rs index 2469fd7..ee0e9e4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -480,6 +480,10 @@ async fn execute_prompt( format!("⚠ Candidate fallback: served by {}", res.model_used).yellow() ); } + // Record budget for each failed attempt plus the successful one + for _ in 0..res.attempts.len() { + let _ = budget_tracker.record_call(false); + } let _ = budget_tracker.record_call(res.tier_used == Tier::Paid); } Err(e) => { @@ -571,6 +575,9 @@ async fn run_pipeline( let _ = event_handler.await; for r in &results { + for _ in 0..r.attempts.len() { + let _ = budget_tracker.record_call(false); + } let _ = budget_tracker.record_call(r.tier_used == Tier::Paid); } diff --git a/src/pipeline.rs b/src/pipeline.rs index e346701..67a153f 100644 --- a/src/pipeline.rs +++ b/src/pipeline.rs @@ -67,13 +67,16 @@ impl PipelineExecutor { let candidates_to_try = if !eligible.is_empty() { eligible - } else { - // If live registry filtering yielded nothing (e.g. offline/strict), fall back to raw persona candidates + } else if model_registry.model_count().await == 0 { + // Uninitialized / offline degraded state: allow raw candidates persona .candidates .iter() .filter(|c| allow_stealth || !c.is_stealth) .collect() + } else { + // Live catalog was fetched and deliberately excluded all candidates as offline/unlisted + vec![] }; if candidates_to_try.is_empty() { @@ -154,9 +157,20 @@ impl PipelineExecutor { _ => { // Continue to next candidate if idx + 1 == total_candidates { + let summary: Vec = attempts + .iter() + .map(|a| { + format!( + "{}: {}", + a.model_id, + a.error.as_deref().unwrap_or("unknown error") + ) + }) + .collect(); return Err(ClientError::ModelUnavailable(format!( - "All {} candidates for persona '{persona_name}' failed. Last error: {err_msg}", - total_candidates + "All {} candidates for persona '{persona_name}' failed:\n {}", + total_candidates, + summary.join("\n ") ))); } } @@ -165,8 +179,19 @@ impl PipelineExecutor { } } + let summary: Vec = attempts + .iter() + .map(|a| { + format!( + "{}: {}", + a.model_id, + a.error.as_deref().unwrap_or("unknown error") + ) + }) + .collect(); Err(ClientError::ModelUnavailable(format!( - "All candidates for persona '{persona_name}' failed." + "All candidates for persona '{persona_name}' failed:\n {}", + summary.join("\n ") ))) } diff --git a/src/router/llm.rs b/src/router/llm.rs index 96ed7c9..d7d8517 100644 --- a/src/router/llm.rs +++ b/src/router/llm.rs @@ -25,6 +25,9 @@ impl LlmRouter { .await { Ok(out) => out, + Err(e @ (ClientError::Unauthorized | ClientError::PaymentRequired)) => { + return Err(e); + } Err(e) => { tracing::warn!("Router LLM call failed, falling back to default persona: {e}"); return Ok(default_persona.to_string()); diff --git a/tests/fallback_tests.rs b/tests/fallback_tests.rs index 324ae87..891dd4b 100644 --- a/tests/fallback_tests.rs +++ b/tests/fallback_tests.rs @@ -172,3 +172,104 @@ async fn test_fail_fast_on_401_unauthorized() { // Must NOT call candidate-2 when 401 Unauthorized occurs! assert_eq!(calls.load(Ordering::SeqCst), 1); } + +#[tokio::test] +async fn test_total_cascade_failure_reports_all_attempts() { + let persona = create_resilience_persona(); + let persona_reg = PersonaRegistry::new(vec![persona]); + let model_reg = ModelRegistry::new(); + + struct AllFailingClient; + #[async_trait] + impl ModelClient for AllFailingClient { + async fn complete( + &self, + model: &str, + _system: &str, + _user: &str, + ) -> Result { + if model == "candidate-1" { + Err(ClientError::ServerError { + status: 503, + message: "Model overloaded".into(), + }) + } else { + Err(ClientError::RateLimited(30)) + } + } + async fn complete_stream( + &self, + model: &str, + system: &str, + user: &str, + _chunk_sender: mpsc::Sender, + ) -> Result { + self.complete(model, system, user).await + } + } + + let err = PipelineExecutor::dispatch_persona( + "resilient-coder", + &persona_reg, + &model_reg, + &AllFailingClient, + "test query", + false, + false, + None, + ) + .await + .unwrap_err(); + + let err_str = err.to_string(); + assert!( + err_str.contains("candidate-1"), + "Error must mention candidate-1: {err_str}" + ); + assert!( + err_str.contains("candidate-2"), + "Error must mention candidate-2: {err_str}" + ); +} + +#[tokio::test] +async fn test_stage_b_router_fails_fast_on_401() { + use brokecrew::router::LlmRouter; + + struct AuthFailingClient; + #[async_trait] + impl ModelClient for AuthFailingClient { + async fn complete( + &self, + _model: &str, + _system: &str, + _user: &str, + ) -> Result { + Err(ClientError::Unauthorized) + } + async fn complete_stream( + &self, + model: &str, + system: &str, + user: &str, + _chunk_sender: mpsc::Sender, + ) -> Result { + self.complete(model, system, user).await + } + } + + let err = LlmRouter::classify( + &AuthFailingClient, + "router-model", + &["coder".into()], + "general", + "ambiguous prompt", + ) + .await + .unwrap_err(); + + match err { + ClientError::Unauthorized => (), + _ => panic!("Expected Unauthorized from router, got {:?}", err), + } +}