fix: address code review findings on cascade failure reporting, auth fail-fast, and budget tracking

This commit is contained in:
katR 2026-09-26 14:04:26 +05:30
parent b7b88c6daa
commit 8c7a08edfc
5 changed files with 153 additions and 7 deletions

View file

@ -194,10 +194,15 @@ impl OpenRouterClient {
temperature, 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 let res = self
.http .http
.post(&endpoint) .post(&endpoint)
.header("Authorization", format!("Bearer {}", self.api_key)) .header("Authorization", auth_val)
.header("HTTP-Referer", "https://github.com/brokecrew/brokecrew") .header("HTTP-Referer", "https://github.com/brokecrew/brokecrew")
.header("X-Title", "BrokeCrew") .header("X-Title", "BrokeCrew")
.header("Content-Type", "application/json") .header("Content-Type", "application/json")
@ -259,10 +264,15 @@ impl ModelClient for OpenRouterClient {
temperature: None, 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 let res = self
.http .http
.post(&endpoint) .post(&endpoint)
.header("Authorization", format!("Bearer {}", self.api_key)) .header("Authorization", auth_val)
.header("HTTP-Referer", "https://github.com/brokecrew/brokecrew") .header("HTTP-Referer", "https://github.com/brokecrew/brokecrew")
.header("X-Title", "BrokeCrew") .header("X-Title", "BrokeCrew")
.header("Content-Type", "application/json") .header("Content-Type", "application/json")

View file

@ -480,6 +480,10 @@ async fn execute_prompt(
format!("⚠ Candidate fallback: served by {}", res.model_used).yellow() 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); let _ = budget_tracker.record_call(res.tier_used == Tier::Paid);
} }
Err(e) => { Err(e) => {
@ -571,6 +575,9 @@ async fn run_pipeline(
let _ = event_handler.await; let _ = event_handler.await;
for r in &results { 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); let _ = budget_tracker.record_call(r.tier_used == Tier::Paid);
} }

View file

@ -67,13 +67,16 @@ impl PipelineExecutor {
let candidates_to_try = if !eligible.is_empty() { let candidates_to_try = if !eligible.is_empty() {
eligible eligible
} else { } else if model_registry.model_count().await == 0 {
// If live registry filtering yielded nothing (e.g. offline/strict), fall back to raw persona candidates // Uninitialized / offline degraded state: allow raw candidates
persona persona
.candidates .candidates
.iter() .iter()
.filter(|c| allow_stealth || !c.is_stealth) .filter(|c| allow_stealth || !c.is_stealth)
.collect() .collect()
} else {
// Live catalog was fetched and deliberately excluded all candidates as offline/unlisted
vec![]
}; };
if candidates_to_try.is_empty() { if candidates_to_try.is_empty() {
@ -154,9 +157,20 @@ impl PipelineExecutor {
_ => { _ => {
// Continue to next candidate // Continue to next candidate
if idx + 1 == total_candidates { if idx + 1 == total_candidates {
let summary: Vec<String> = attempts
.iter()
.map(|a| {
format!(
"{}: {}",
a.model_id,
a.error.as_deref().unwrap_or("unknown error")
)
})
.collect();
return Err(ClientError::ModelUnavailable(format!( return Err(ClientError::ModelUnavailable(format!(
"All {} candidates for persona '{persona_name}' failed. Last error: {err_msg}", "All {} candidates for persona '{persona_name}' failed:\n {}",
total_candidates total_candidates,
summary.join("\n ")
))); )));
} }
} }
@ -165,8 +179,19 @@ impl PipelineExecutor {
} }
} }
let summary: Vec<String> = attempts
.iter()
.map(|a| {
format!(
"{}: {}",
a.model_id,
a.error.as_deref().unwrap_or("unknown error")
)
})
.collect();
Err(ClientError::ModelUnavailable(format!( Err(ClientError::ModelUnavailable(format!(
"All candidates for persona '{persona_name}' failed." "All candidates for persona '{persona_name}' failed:\n {}",
summary.join("\n ")
))) )))
} }

View file

@ -25,6 +25,9 @@ impl LlmRouter {
.await .await
{ {
Ok(out) => out, Ok(out) => out,
Err(e @ (ClientError::Unauthorized | ClientError::PaymentRequired)) => {
return Err(e);
}
Err(e) => { Err(e) => {
tracing::warn!("Router LLM call failed, falling back to default persona: {e}"); tracing::warn!("Router LLM call failed, falling back to default persona: {e}");
return Ok(default_persona.to_string()); return Ok(default_persona.to_string());

View file

@ -172,3 +172,104 @@ async fn test_fail_fast_on_401_unauthorized() {
// Must NOT call candidate-2 when 401 Unauthorized occurs! // Must NOT call candidate-2 when 401 Unauthorized occurs!
assert_eq!(calls.load(Ordering::SeqCst), 1); 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<String, ClientError> {
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<String>,
) -> Result<String, ClientError> {
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<String, ClientError> {
Err(ClientError::Unauthorized)
}
async fn complete_stream(
&self,
model: &str,
system: &str,
user: &str,
_chunk_sender: mpsc::Sender<String>,
) -> Result<String, ClientError> {
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),
}
}