From 9fb0874b61308160a0ab3e3d6f00383501717ab1 Mon Sep 17 00:00:00 2001 From: oiupoyt Date: Sun, 11 Oct 2026 18:26:50 +0530 Subject: [PATCH] fix: architectural review items 1-10 (session ids, threshold wiring, utf8 truncation, pipeline resilience, truthful telemetry, registry sync, arcswap config, dispatch params, atomic budget, total check_status) --- Cargo.toml | 2 + src/budget.rs | 74 ++++++- src/client.rs | 83 ++++---- src/config.rs | 32 +++ src/lib.rs | 5 +- src/main.rs | 98 +++++---- src/pipeline.rs | 431 ++++++++++++++++++++++++++++++++-------- src/registry.rs | 34 +++- src/router/mod.rs | 43 ++-- src/session.rs | 15 +- src/tool.rs | 19 +- src/tui/app.rs | 59 +----- src/tui/mod.rs | 196 ++++++++++-------- src/tui/ui.rs | 37 ---- tests/fallback_tests.rs | 82 ++++---- 15 files changed, 805 insertions(+), 405 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 2e356e4..ea61011 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,6 +28,8 @@ owo-colors = "4.1" regex = "1.10" chrono = { version = "0.4", features = ["serde"] } dirs = "5.0" +rand = "0.8" +arc-swap = "1.7" tempfile = "3.12" ratatui = "0.29" crossterm = { version = "0.28", features = ["event-stream"] } diff --git a/src/budget.rs b/src/budget.rs index b465cdf..5f46e68 100644 --- a/src/budget.rs +++ b/src/budget.rs @@ -36,6 +36,7 @@ pub struct BudgetStatus { pub minute_limit: usize, pub daily_requests: usize, pub daily_limit: usize, + pub daily_paid_requests: usize, pub approaching_daily_limit: bool, } @@ -50,6 +51,7 @@ pub enum BudgetExceeded { #[derive(Debug, Serialize, Deserialize, Default)] struct StoredBudgetState { daily_count: usize, + daily_paid_count: usize, last_day_utc: String, // YYYY-MM-DD recent_calls: VecDeque>, } @@ -58,6 +60,7 @@ pub struct BudgetTracker { storage_path: Option, credit_tier: CreditTier, daily_count: usize, + daily_paid_count: usize, last_day_utc: String, recent_calls: VecDeque>, } @@ -68,6 +71,7 @@ impl BudgetTracker { storage_path, credit_tier, daily_count: 0, + daily_paid_count: 0, last_day_utc: Utc::now().format("%Y-%m-%d").to_string(), recent_calls: VecDeque::new(), }; @@ -91,6 +95,7 @@ impl BudgetTracker { // Day rollover check if self.last_day_utc != current_day { self.daily_count = 0; + self.daily_paid_count = 0; self.last_day_utc = current_day; } @@ -132,16 +137,49 @@ impl BudgetTracker { minute_limit: MINUTE_LIMIT, daily_requests: self.daily_count, daily_limit, + daily_paid_requests: self.daily_paid_count, approaching_daily_limit, }) } + /// Atomically checks limits and records one call. Replaces the racy + /// check-then-record sequence: a single call either acquires budget or + /// reports exhaustion, never both passing then both recording. + pub fn try_acquire(&mut self, is_paid: bool) -> Result { + // Reuse check_budget for limit evaluation (it refreshes the window). + let _ = self.check_budget()?; + let now = Utc::now(); + self.recent_calls.push_back(now); + self.daily_count += 1; + if is_paid { + self.daily_paid_count += 1; + } + + if let Some(ref path) = self.storage_path.clone() { + let _ = self.save_to_disk(path); + } + + // Status after recording. + let daily_limit = self.credit_tier.daily_limit(); + Ok(BudgetStatus { + minute_requests: self.recent_calls.len(), + minute_limit: MINUTE_LIMIT, + daily_requests: self.daily_count, + daily_limit, + daily_paid_requests: self.daily_paid_count, + approaching_daily_limit: self.daily_count >= (daily_limit * 80) / 100, + }) + } + /// Records an outbound call and persists updated state to disk atomically. - pub fn record_call(&mut self, _is_paid: bool) -> Result<(), std::io::Error> { + pub fn record_call(&mut self, is_paid: bool) -> Result<(), std::io::Error> { self.refresh_window(); let now = Utc::now(); self.recent_calls.push_back(now); self.daily_count += 1; + if is_paid { + self.daily_paid_count += 1; + } if let Some(ref path) = self.storage_path.clone() { self.save_to_disk(path)?; @@ -156,6 +194,7 @@ impl BudgetTracker { let data = fs::read_to_string(path)?; if let Ok(state) = serde_json::from_str::(&data) { self.daily_count = state.daily_count; + self.daily_paid_count = state.daily_paid_count; self.last_day_utc = state.last_day_utc; self.recent_calls = state.recent_calls; } @@ -168,6 +207,7 @@ impl BudgetTracker { } let state = StoredBudgetState { daily_count: self.daily_count, + daily_paid_count: self.daily_paid_count, last_day_utc: self.last_day_utc.clone(), recent_calls: self.recent_calls.clone(), }; @@ -245,4 +285,36 @@ mod tests { assert_eq!(status.daily_requests, 2); assert_eq!(status.minute_requests, 2); } + + #[test] + fn test_try_acquire_tracks_paid_and_enforces_limits() { + let mut tracker = BudgetTracker::new(None, CreditTier::LowBalance); + let s1 = tracker.try_acquire(false).expect("first acquire"); + assert_eq!(s1.daily_requests, 1); + assert_eq!(s1.daily_paid_requests, 0); + let s2 = tracker.try_acquire(true).expect("paid acquire"); + assert_eq!(s2.daily_requests, 2); + assert_eq!(s2.daily_paid_requests, 1); + + // Fill the minute window through try_acquire: 18 more on top of the + // 2 already recorded. At 20 calls the window is full, so check_budget + // reports the limit (it returns Err at capacity, not Ok(20)). + for _ in 0..18 { + tracker.try_acquire(false).expect("window fill"); + } + match tracker.check_budget() { + Err(BudgetExceeded::MinuteLimit(20, _)) => (), + other => panic!("window should be exactly full, got {other:?}"), + } + // The next acquire must fail atomically (no record on failure). + // NOTE: wait-seconds hints differ per call, so match variants only. + match tracker.try_acquire(false) { + Err(BudgetExceeded::MinuteLimit(_, _)) => (), + other => panic!("expected minute limit, got {other:?}"), + } + match tracker.check_budget() { + Err(BudgetExceeded::MinuteLimit(_, _)) => (), + other => panic!("expected still-limited window, got {other:?}"), + } + } } diff --git a/src/client.rs b/src/client.rs index 2ce59f3..db5da34 100644 --- a/src/client.rs +++ b/src/client.rs @@ -3,7 +3,7 @@ use futures_util::StreamExt; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::fmt; -use std::sync::RwLock; +use std::sync::Arc; use std::time::Duration; use tokio::sync::mpsc; @@ -38,18 +38,25 @@ pub trait ModelClient: Send + Sync { ) -> Result; } +/// Immutable snapshot of client credentials. Rebuilt wholesale on +/// /provider or /key change and swapped behind an Arc, so request paths +/// never block on a lock. +#[derive(Debug, Clone)] +struct ClientConfig { + api_key: String, + base_url: String, +} + pub struct OpenRouterClient { - api_key: RwLock, - base_url: RwLock, + config: Arc>, http: reqwest::Client, } impl fmt::Debug for OpenRouterClient { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let base_url = self.base_url.read().map(|b| b.clone()).unwrap_or_default(); f.debug_struct("OpenRouterClient") .field("api_key", &"[REDACTED]") - .field("base_url", &base_url) + .field("base_url", &self.config.load().base_url) .finish() } } @@ -99,30 +106,36 @@ impl OpenRouterClient { .unwrap_or_else(|_| reqwest::Client::new()); Self { - api_key: RwLock::new(api_key), - base_url: RwLock::new(base_url), + config: Arc::new(arc_swap::ArcSwap::new(Arc::new(ClientConfig { + api_key, + base_url, + }))), http, } } pub fn set_api_key(&self, key: String) { - if let Ok(mut write) = self.api_key.write() { - *write = key; - } + let current = self.config.load(); + self.config.store(Arc::new(ClientConfig { + api_key: key, + base_url: current.base_url.clone(), + })); } pub fn set_base_url(&self, url: String) { - if let Ok(mut write) = self.base_url.write() { - *write = url; - } + let current = self.config.load(); + self.config.store(Arc::new(ClientConfig { + api_key: current.api_key.clone(), + base_url: url, + })); } pub fn get_api_key(&self) -> String { - self.api_key.read().map(|k| k.clone()).unwrap_or_default() + self.config.load().api_key.clone() } pub fn get_base_url(&self) -> String { - self.base_url.read().map(|b| b.clone()).unwrap_or_default() + self.config.load().base_url.clone() } /// Parses an SSE formatted line into a token delta if available @@ -158,31 +171,30 @@ impl OpenRouterClient { None } + /// Maps an HTTP response to an optional client error. `None` means success. + /// Returns `Option` instead of `Result` so error paths stay total: there + /// is no `Ok` value to unwrap out from under a future caller. fn check_status( status: reqwest::StatusCode, body: &str, headers: &reqwest::header::HeaderMap, - ) -> Result<(), ClientError> { + ) -> Option { match status.as_u16() { - 200..=299 => Ok(()), - 401 => Err(ClientError::Unauthorized), - 402 => Err(ClientError::PaymentRequired), + 200..=299 => None, + 401 => Some(ClientError::Unauthorized), + 402 => Some(ClientError::PaymentRequired), 429 => { let retry_after = headers .get("retry-after") .and_then(|h| h.to_str().ok()) .and_then(|s| s.parse::().ok()) .unwrap_or(5); - Err(ClientError::RateLimited(retry_after)) + Some(ClientError::RateLimited(retry_after)) } - 404 => Err(ClientError::ModelUnavailable(format!( + 404 => Some(ClientError::ModelUnavailable(format!( "Model not found (404): {body}" ))), - s if s >= 500 => Err(ClientError::ServerError { - status: s, - message: body.to_string(), - }), - s => Err(ClientError::ServerError { + s => Some(ClientError::ServerError { status: s, message: body.to_string(), }), @@ -236,9 +248,10 @@ impl OpenRouterClient { let status = res.status(); let headers = res.headers().clone(); - if !status.is_success() { + if let Some(err) = Self::check_status(status, "", &headers) { let error_text = res.text().await.unwrap_or_default(); - return Err(Self::check_status(status, &error_text, &headers).unwrap_err()); + // Re-map with the body so 4xx/5xx messages carry server detail. + return Err(Self::check_status(status, &error_text, &headers).unwrap_or(err)); } let resp_json: ChatCompletionResponse = res.json().await?; @@ -307,9 +320,9 @@ impl ModelClient for OpenRouterClient { let status = res.status(); let headers = res.headers().clone(); - if !status.is_success() { + if let Some(err) = Self::check_status(status, "", &headers) { let error_text = res.text().await.unwrap_or_default(); - return Err(Self::check_status(status, &error_text, &headers).unwrap_err()); + return Err(Self::check_status(status, &error_text, &headers).unwrap_or(err)); } let mut stream = res.bytes_stream(); @@ -377,9 +390,11 @@ mod tests { reqwest::header::HeaderValue::from_static("12"), ); + assert!(OpenRouterClient::check_status(reqwest::StatusCode::OK, "", &headers).is_none()); + let err_401 = OpenRouterClient::check_status(reqwest::StatusCode::UNAUTHORIZED, "bad key", &headers) - .unwrap_err(); + .expect("401 must map to an error"); match err_401 { ClientError::Unauthorized => (), _ => panic!("Expected Unauthorized, got {:?}", err_401), @@ -390,7 +405,7 @@ mod tests { "low balance", &headers, ) - .unwrap_err(); + .expect("402 must map to an error"); match err_402 { ClientError::PaymentRequired => (), _ => panic!("Expected PaymentRequired, got {:?}", err_402), @@ -401,7 +416,7 @@ mod tests { "rate limited", &headers, ) - .unwrap_err(); + .expect("429 must map to an error"); match err_429 { ClientError::RateLimited(secs) => assert_eq!(secs, 12), _ => panic!("Expected RateLimited(12), got {:?}", err_429), @@ -412,7 +427,7 @@ mod tests { "outage", &headers, ) - .unwrap_err(); + .expect("503 must map to an error"); 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 cd1e9dd..963e939 100644 --- a/src/config.rs +++ b/src/config.rs @@ -212,6 +212,28 @@ impl AppConfig { Ok(()) } + /// Persists the complexity escalation threshold into `config.toml` + pub fn save_complexity_threshold(config_dir: &Path, threshold: u8) -> Result<(), String> { + let config_file_path = config_dir.join("config.toml"); + let mut 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).unwrap_or_default() + } else { + AppConfigFile::default() + }; + + app_file.complexity_threshold = threshold.clamp(1, 10); + let serialized = toml::to_string_pretty(&app_file) + .map_err(|e| format!("Failed to serialize config: {e}"))?; + + fs::create_dir_all(config_dir).map_err(|e| format!("Failed to create config dir: {e}"))?; + fs::write(&config_file_path, serialized) + .map_err(|e| format!("Failed to write {}: {e}", config_file_path.display()))?; + + Ok(()) + } + /// Masks an API key for safe display (e.g. `sk-or-v••••••••`) pub fn mask_api_key(key: &str) -> String { let trimmed = key.trim(); @@ -359,4 +381,14 @@ mod tests { "sk-or-v••••••••" ); } + + #[test] + fn test_complexity_threshold_persistence() { + let temp_dir = tempfile::tempdir().unwrap(); + AppConfig::save_complexity_threshold(temp_dir.path(), 9).expect("save threshold"); + let loaded = + AppConfig::load_from_paths(Some(&temp_dir.path().join("config.toml")), None, None) + .expect("load config"); + assert_eq!(loaded.complexity_threshold, 9); + } } diff --git a/src/lib.rs b/src/lib.rs index 5817367..d8dfdb5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,7 +14,10 @@ pub use budget::{estimate_tokens, BudgetExceeded, BudgetStatus, BudgetTracker, C 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 pipeline::{ + CliRunParams, DispatchParams, PipelineConfig, PipelineEvent, PipelineExecutor, PipelineOutcome, + PipelineRunParams, PipelineStep, +}; pub use registry::{ModelInfo, ModelPricing, ModelRegistry}; pub use router::{KeywordRouter, LlmRouter, RouteDecision, RouteStage, Router}; pub use session::{Session, SessionMessage, SessionSummary}; diff --git a/src/main.rs b/src/main.rs index 635205c..9032458 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,7 +5,9 @@ use brokecrew::budget::BudgetTracker; use brokecrew::client::OpenRouterClient; use brokecrew::config::AppConfig; use brokecrew::persona::Tier; -use brokecrew::pipeline::{PipelineConfig, PipelineEvent, PipelineExecutor}; +use brokecrew::pipeline::{ + CliRunParams, DispatchParams, PipelineEvent, PipelineExecutor, PipelineRunParams, +}; use brokecrew::registry::ModelRegistry; use brokecrew::router::{RouteStage, Router}; use clap::Parser; @@ -422,12 +424,11 @@ async fn run_chat( let client = OpenRouterClient::with_base_url(api_key, config.openrouter_base_url.clone()); - // Start background sync for model registry - let reg_clone = model_registry.clone(); - let base_url = config.openrouter_base_url.clone(); - tokio::spawn(async move { - let _ = reg_clone.fetch_live(&base_url).await; - }); + // Start background sync for model registry (immediate + every 30 min). + let _registry_sync = model_registry.start_background_sync( + config.openrouter_base_url.clone(), + std::time::Duration::from_secs(30 * 60), + ); let router = Router::new( config.default_persona.clone(), @@ -505,7 +506,7 @@ async fn execute_prompt( router: &Router, client: &OpenRouterClient, ) -> Result<(), Box> { - // 1. Budget check + // 1. Budget pre-check (advisory only; each dispatch records atomically). match budget_tracker.check_budget() { Ok(status) => { if status.approaching_daily_limit { @@ -534,7 +535,7 @@ async fn execute_prompt( .find(|p| p.name == *pipe_name) .ok_or_else(|| format!("Pipeline '{pipe_name}' not found in configuration"))?; - run_pipeline( + run_pipeline(CliRunParams { pipeline, config, budget_tracker, @@ -543,7 +544,7 @@ async fn execute_prompt( prompt, allow_stealth, require_paid, - ) + }) .await?; return Ok(()); } @@ -562,7 +563,7 @@ async fn execute_prompt( match decision { Ok(dec) => { if dec.stage_used == RouteStage::StageBLlm { - let _ = budget_tracker.record_call(false); + let _ = budget_tracker.try_acquire(false); } let info = match dec.stage_used { RouteStage::StageAKeyword => "routed via keyword pre-filter".to_string(), @@ -591,16 +592,16 @@ async fn execute_prompt( let persona_tag = get_persona_color_tag(&target_persona); let spinner = create_spinner("Thinking..."); - let dispatch_future = PipelineExecutor::dispatch_persona( - &target_persona, - &config.persona_registry, + let dispatch_future = PipelineExecutor::dispatch_persona(DispatchParams { + persona_name: &target_persona, + persona_registry: &config.persona_registry, model_registry, client, prompt, allow_stealth, require_paid, - Some(chunk_tx), - ); + chunk_sender: Some(chunk_tx), + }); let printer = tokio::spawn(async move { let mut first = true; @@ -629,9 +630,9 @@ async fn execute_prompt( } // 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.try_acquire(false); } - let _ = budget_tracker.record_call(res.tier_used == Tier::Paid); + let _ = budget_tracker.try_acquire(res.tier_used == Tier::Paid); } Err(e) => { eprintln!("\n{} {}", "Error dispatching to persona:".red().bold(), e); @@ -641,17 +642,17 @@ async fn execute_prompt( Ok(()) } -#[allow(clippy::too_many_arguments)] -async fn run_pipeline( - pipeline: &PipelineConfig, - config: &AppConfig, - budget_tracker: &mut BudgetTracker, - model_registry: &ModelRegistry, - client: &OpenRouterClient, - prompt: &str, - allow_stealth: bool, - require_paid: bool, -) -> Result<(), Box> { +async fn run_pipeline(params: CliRunParams<'_>) -> Result<(), Box> { + let CliRunParams { + pipeline, + config, + budget_tracker, + model_registry, + client, + prompt, + allow_stealth, + require_paid, + } = params; println!( "\n{} Starting pipeline '{}' ({} steps)", "›".bright_cyan(), @@ -661,16 +662,16 @@ async fn run_pipeline( let (event_tx, mut event_rx) = mpsc::channel::(100); - let exec_future = PipelineExecutor::execute( + let exec_future = PipelineExecutor::execute(PipelineRunParams { pipeline, - &config.persona_registry, + persona_registry: &config.persona_registry, model_registry, client, - prompt, + initial_prompt: prompt, allow_stealth, require_paid, - Some(event_tx), - ); + event_sender: Some(event_tx), + }); let event_handler = tokio::spawn(async move { let mut current_spinner: Option = None; @@ -681,23 +682,33 @@ async fn run_pipeline( step_index, total_steps, persona, - model, - .. } => { if let Some(sp) = current_spinner.take() { sp.finish_and_clear(); } let badge = get_persona_color_tag(&persona); println!( - "\n--- Step {}/{} {} ({}) ---", + "\n--- Step {}/{} {} ---", step_index + 1, total_steps, badge, - model.dimmed() ); let sp = create_spinner("Thinking..."); current_spinner = Some(sp); } + PipelineEvent::StepFailed { + step_index, + persona, + error, + } => { + if let Some(sp) = current_spinner.take() { + sp.finish_and_clear(); + } + eprintln!( + "{}", + format!("! Step {}/{persona} failed: {error}", step_index + 1,).red() + ); + } PipelineEvent::CandidateAttempt { model, attempt, @@ -718,11 +729,16 @@ async fn run_pipeline( print!("{delta}"); let _ = io::stdout().flush(); } - PipelineEvent::StepCompleted { fell_back, .. } => { + PipelineEvent::StepCompleted { + model_used, + fell_back, + .. + } => { if let Some(sp) = current_spinner.take() { sp.finish_and_clear(); } println!(); + eprintln!("{}", format!("› served by {model_used}").dimmed()); if fell_back { eprintln!("{}", "! Fallback occurred during step execution".yellow()); } @@ -736,9 +752,9 @@ async fn run_pipeline( for r in &results { for _ in 0..r.attempts.len() { - let _ = budget_tracker.record_call(false); + let _ = budget_tracker.try_acquire(false); } - let _ = budget_tracker.record_call(r.tier_used == Tier::Paid); + let _ = budget_tracker.try_acquire(r.tier_used == Tier::Paid); } Ok(()) diff --git a/src/pipeline.rs b/src/pipeline.rs index bdf646d..2fd0f32 100644 --- a/src/pipeline.rs +++ b/src/pipeline.rs @@ -1,4 +1,6 @@ -use crate::client::{ClientError, ModelClient}; +use crate::budget::BudgetTracker; +use crate::client::{ClientError, ModelClient, OpenRouterClient}; +use crate::config::AppConfig; use crate::persona::{DispatchResult, ModelAttempt, ModelCandidate, PersonaRegistry, Tier}; use crate::registry::ModelRegistry; use serde::{Deserialize, Serialize}; @@ -17,6 +19,23 @@ pub struct PipelineStep { pub struct PipelineConfig { pub name: String, pub steps: Vec, + /// When false (default), a failed step aborts the pipeline but completed + /// step results are preserved in the returned outcome. + #[serde(default = "default_abort_on_failure")] + pub abort_on_failure: bool, +} + +fn default_abort_on_failure() -> bool { + true +} + +/// Outcome of a pipeline run: completed step results plus, on failure, +/// which step failed and why. Callers keep partial results either way. +#[derive(Debug)] +pub struct PipelineOutcome { + pub results: Vec, + pub failed_step: Option, + pub error: Option, } #[derive(Debug, Clone)] @@ -25,8 +44,6 @@ pub enum PipelineEvent { step_index: usize, total_steps: usize, persona: String, - model: String, - tier: Tier, }, CandidateAttempt { step_index: usize, @@ -49,28 +66,74 @@ pub enum PipelineEvent { StepCompleted { step_index: usize, persona: String, + model_used: String, + tier_used: Tier, full_reply: String, fell_back: bool, attempts: Vec, }, + StepFailed { + step_index: usize, + persona: String, + error: String, + }, +} + +/// Parameters for a single persona dispatch. Bundles the 8-argument +/// call sites into one value so callers stay readable. +#[derive(Clone)] +pub struct DispatchParams<'a> { + pub persona_name: &'a str, + pub persona_registry: &'a PersonaRegistry, + pub model_registry: &'a ModelRegistry, + pub client: &'a dyn ModelClient, + pub prompt: &'a str, + pub allow_stealth: bool, + pub require_paid: bool, + pub chunk_sender: Option>, +} + +/// Parameters for a full pipeline run. See [`DispatchParams`]. +#[derive(Clone)] +pub struct PipelineRunParams<'a> { + pub pipeline: &'a PipelineConfig, + pub persona_registry: &'a PersonaRegistry, + pub model_registry: &'a ModelRegistry, + pub client: &'a dyn ModelClient, + pub initial_prompt: &'a str, + pub allow_stealth: bool, + pub require_paid: bool, + pub event_sender: Option>, +} + +/// Parameters for the CLI `run_pipeline` presentation loop. +pub struct CliRunParams<'a> { + pub pipeline: &'a PipelineConfig, + pub config: &'a AppConfig, + pub budget_tracker: &'a mut BudgetTracker, + pub model_registry: &'a ModelRegistry, + pub client: &'a OpenRouterClient, + pub prompt: &'a str, + pub allow_stealth: bool, + pub require_paid: bool, } 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, - model_registry: &ModelRegistry, - client: &dyn ModelClient, - prompt: &str, - allow_stealth: bool, - require_paid: bool, - chunk_sender: Option>, + params: DispatchParams<'_>, ) -> Result { - Self::dispatch_persona_with_events( + Self::dispatch_persona_with_events(params, None).await + } + + /// Executes a single persona query with candidate fallback walk and telemetry event streaming + pub async fn dispatch_persona_with_events( + params: DispatchParams<'_>, + event_sender: Option>, + ) -> Result { + let DispatchParams { persona_name, persona_registry, model_registry, @@ -79,24 +142,7 @@ impl PipelineExecutor { allow_stealth, require_paid, chunk_sender, - None, - ) - .await - } - - /// Executes a single persona query with candidate fallback walk and telemetry event streaming - #[allow(clippy::too_many_arguments)] - pub async fn dispatch_persona_with_events( - persona_name: &str, - persona_registry: &PersonaRegistry, - model_registry: &ModelRegistry, - client: &dyn ModelClient, - prompt: &str, - allow_stealth: bool, - require_paid: bool, - chunk_sender: Option>, - event_sender: Option>, - ) -> Result { + } = params; let persona = persona_registry.get(persona_name).ok_or_else(|| { ClientError::ModelUnavailable(format!("Persona '{persona_name}' not found")) })?; @@ -274,26 +320,55 @@ impl PipelineExecutor { ))) } - /// Executes an entire pipeline sequentially, passing outputs between steps - #[allow(clippy::too_many_arguments)] + /// Executes an entire pipeline sequentially, passing outputs between steps. + /// On step failure the outcome carries completed results plus the failure; + /// callers decide whether partial output is usable. `abort_on_failure=false` + /// continues with the remaining steps instead of stopping. pub async fn execute( - pipeline: &PipelineConfig, - persona_registry: &PersonaRegistry, - model_registry: &ModelRegistry, - client: &dyn ModelClient, - initial_prompt: &str, - allow_stealth: bool, - require_paid: bool, - event_sender: Option>, + params: PipelineRunParams<'_>, ) -> Result, ClientError> { + let outcome = Self::execute_outcome(params).await; + match outcome.error { + None => Ok(outcome.results), + Some(e) => Err(e), + } + } + + /// Variant of [`Self::execute`] that preserves partial results on failure. + pub async fn execute_outcome(params: PipelineRunParams<'_>) -> PipelineOutcome { + let PipelineRunParams { + pipeline, + persona_registry, + model_registry, + client, + initial_prompt, + allow_stealth, + require_paid, + event_sender, + } = params; let mut results = Vec::new(); let mut previous_output = String::new(); 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)) - })?; + if persona_registry.get(&step.persona).is_none() { + let err = + ClientError::ModelUnavailable(format!("Persona '{}' not found", step.persona)); + if let Some(ref tx) = event_sender { + let _ = tx + .send(PipelineEvent::StepFailed { + step_index: step_idx, + persona: step.persona.clone(), + error: err.to_string(), + }) + .await; + } + return PipelineOutcome { + results, + failed_step: Some(step_idx), + error: Some(err), + }; + } // Prepare prompt for this step let step_prompt = if step.pass_previous_output && !previous_output.is_empty() { @@ -327,59 +402,84 @@ impl PipelineExecutor { } }); - // Notify step start + // Notify step start. The model is intentionally omitted here: + // the fallback walk has not run yet, so any advertised model + // would be a guess. The resolved model arrives via + // CandidateAttempt events and the StepCompleted model_used field. if let Some(ref tx) = event_sender { - // Peek top candidate to inform UI - let top_candidate = persona - .candidates - .first() - .map(|c| (c.id.clone(), c.tier)) - .unwrap_or_else(|| ("unknown".into(), Tier::Free)); - let _ = tx .send(PipelineEvent::StepStarted { step_index: step_idx, total_steps, persona: step.persona.clone(), - model: top_candidate.0, - tier: top_candidate.1, }) .await; } let dispatch_result = Self::dispatch_persona_with_events( - &step.persona, - persona_registry, - model_registry, - client, - &step_prompt, - allow_stealth, - require_paid, - Some(chunk_tx), + DispatchParams { + persona_name: &step.persona, + persona_registry, + model_registry, + client, + prompt: &step_prompt, + allow_stealth, + require_paid, + chunk_sender: Some(chunk_tx), + }, event_sender.clone(), ) - .await?; + .await; let _ = forwarder.await; - previous_output = dispatch_result.reply.clone(); + match dispatch_result { + Ok(dispatch_result) => { + previous_output = dispatch_result.reply.clone(); - if let Some(ref tx) = event_sender { - let _ = tx - .send(PipelineEvent::StepCompleted { - step_index: step_idx, - persona: step.persona.clone(), - full_reply: dispatch_result.reply.clone(), - fell_back: dispatch_result.fell_back, - attempts: dispatch_result.attempts.clone(), - }) - .await; + if let Some(ref tx) = event_sender { + let _ = tx + .send(PipelineEvent::StepCompleted { + step_index: step_idx, + persona: step.persona.clone(), + model_used: dispatch_result.model_used.clone(), + tier_used: dispatch_result.tier_used, + full_reply: dispatch_result.reply.clone(), + fell_back: dispatch_result.fell_back, + attempts: dispatch_result.attempts.clone(), + }) + .await; + } + + results.push(dispatch_result); + } + Err(e) => { + if let Some(ref tx) = event_sender { + let _ = tx + .send(PipelineEvent::StepFailed { + step_index: step_idx, + persona: step.persona.clone(), + error: e.to_string(), + }) + .await; + } + if pipeline.abort_on_failure { + return PipelineOutcome { + results, + failed_step: Some(step_idx), + error: Some(e), + }; + } + // Continue: keep previous_output from the last success. + } } - - results.push(dispatch_result); } - Ok(results) + PipelineOutcome { + results, + failed_step: None, + error: None, + } } } @@ -461,6 +561,7 @@ mod tests { custom_prompt_template: None, }, ], + abort_on_failure: true, }; let calls = Arc::new(AtomicUsize::new(0)); @@ -468,16 +569,16 @@ mod tests { calls: calls.clone(), }; - let results = PipelineExecutor::execute( - &pipeline, - &persona_reg, - &model_reg, - &client, - "Build website", - false, - false, - None, - ) + let results = PipelineExecutor::execute(PipelineRunParams { + pipeline: &pipeline, + persona_registry: &persona_reg, + model_registry: &model_reg, + client: &client, + initial_prompt: "Build website", + allow_stealth: false, + require_paid: false, + event_sender: None, + }) .await .expect("pipeline execution failed"); @@ -490,4 +591,160 @@ mod tests { assert!(results[1].reply.contains("Output from coder-model")); assert_eq!(calls.load(Ordering::SeqCst), 2); } + + struct AlwaysFailClient; + + #[async_trait] + impl ModelClient for AlwaysFailClient { + async fn complete( + &self, + _model: &str, + _system: &str, + _user: &str, + ) -> Result { + Err(ClientError::ServerError { + status: 503, + message: "down".into(), + }) + } + + async fn complete_stream( + &self, + model: &str, + system: &str, + user: &str, + _chunk_sender: mpsc::Sender, + ) -> Result { + self.complete(model, system, user).await + } + } + + fn failing_two_step_pipeline() -> (PersonaRegistry, PipelineConfig) { + let mk = |name: &str| Persona { + name: name.into(), + system_prompt: name.into(), + trigger_words: vec![], + candidates: vec![ModelCandidate { + id: format!("{name}-model"), + tier: Tier::Free, + is_stealth: false, + }], + }; + let reg = PersonaRegistry::new(vec![mk("coder"), mk("reviewer")]); + let pipe = PipelineConfig { + name: "fail-fast".into(), + steps: vec![ + PipelineStep { + persona: "coder".into(), + pass_previous_output: false, + custom_prompt_template: None, + }, + PipelineStep { + persona: "reviewer".into(), + pass_previous_output: true, + custom_prompt_template: None, + }, + ], + abort_on_failure: true, + }; + (reg, pipe) + } + + #[tokio::test] + async fn test_pipeline_failure_preserves_partial_results() { + struct FailSecondClient; + #[async_trait] + impl ModelClient for FailSecondClient { + async fn complete( + &self, + model: &str, + _system: &str, + user: &str, + ) -> Result { + if model == "reviewer-model" { + return Err(ClientError::ServerError { + status: 503, + message: "reviewer down".into(), + }); + } + Ok(format!("Output from {model} for: {user}")) + } + async fn complete_stream( + &self, + model: &str, + system: &str, + user: &str, + chunk_sender: mpsc::Sender, + ) -> Result { + let res = self.complete(model, system, user).await?; + let _ = chunk_sender.send(res.clone()).await; + Ok(res) + } + } + + let (reg, pipe) = failing_two_step_pipeline(); + let model_reg = ModelRegistry::new(); + let (event_tx, mut event_rx) = mpsc::channel::(16); + let outcome = PipelineExecutor::execute_outcome(PipelineRunParams { + pipeline: &pipe, + persona_registry: ®, + model_registry: &model_reg, + client: &FailSecondClient, + initial_prompt: "Build website", + allow_stealth: false, + require_paid: false, + event_sender: Some(event_tx), + }) + .await; + + // First step succeeded and is preserved despite the abort. + assert_eq!(outcome.results.len(), 1); + assert_eq!(outcome.failed_step, Some(1)); + assert!(outcome.error.is_some()); + + // A StepFailed event must have been emitted for the UI. + let mut saw_failed = false; + while let Ok(ev) = event_rx.try_recv() { + if matches!(ev, PipelineEvent::StepFailed { step_index: 1, .. }) { + saw_failed = true; + } + } + assert!(saw_failed, "expected StepFailed event for step 1"); + + // Legacy execute() still surfaces the error. + let err = PipelineExecutor::execute(PipelineRunParams { + pipeline: &pipe, + persona_registry: ®, + model_registry: &model_reg, + client: &FailSecondClient, + initial_prompt: "Build website", + allow_stealth: false, + require_paid: false, + event_sender: None, + }) + .await + .unwrap_err(); + assert!(err.to_string().contains("reviewer down")); + } + + #[tokio::test] + async fn test_pipeline_continue_on_failure_policy() { + let (reg, mut pipe) = failing_two_step_pipeline(); + pipe.abort_on_failure = false; + let model_reg = ModelRegistry::new(); + let outcome = PipelineExecutor::execute_outcome(PipelineRunParams { + pipeline: &pipe, + persona_registry: ®, + model_registry: &model_reg, + client: &AlwaysFailClient, + initial_prompt: "Build website", + allow_stealth: false, + require_paid: false, + event_sender: None, + }) + .await; + assert_eq!(outcome.failed_step, None); + assert!(outcome.error.is_none()); + assert!(outcome.results.is_empty()); + } } diff --git a/src/registry.rs b/src/registry.rs index 21ad9cc..031ffea 100644 --- a/src/registry.rs +++ b/src/registry.rs @@ -33,6 +33,7 @@ struct OpenRouterModelsResponse { pub struct ModelRegistry { models: Arc>>, last_refreshed: Arc>>, + http: reqwest::Client, } impl Default for ModelRegistry { @@ -43,9 +44,14 @@ impl Default for ModelRegistry { impl ModelRegistry { pub fn new() -> Self { + let http = reqwest::Client::builder() + .timeout(Duration::from_secs(15)) + .build() + .unwrap_or_else(|_| reqwest::Client::new()); Self { models: Arc::new(RwLock::new(HashMap::new())), last_refreshed: Arc::new(RwLock::new(None)), + http, } } @@ -84,12 +90,9 @@ impl ModelRegistry { /// Fetches live models catalog from OpenRouter `/v1/models` pub async fn fetch_live(&self, base_url: &str) -> Result { let endpoint = format!("{}/models", base_url.trim_end_matches('/')); - let client = reqwest::Client::builder() - .timeout(Duration::from_secs(15)) - .build() - .map_err(|e| e.to_string())?; - let res = client + let res = self + .http .get(&endpoint) .header("HTTP-Referer", "https://github.com/brokecrew/brokecrew") .header("X-Title", "BrokeCrew") @@ -112,19 +115,32 @@ impl ModelRegistry { self.update_from_json(&body).await } - /// Starts a background task refreshing the models registry periodically - pub fn start_background_sync(self: Arc, base_url: String, interval: Duration) { + /// Starts a background task refreshing the models registry periodically. + /// Returns a handle the caller can abort on shutdown; the task does an + /// immediate fetch, then repeats on `interval` (30 min recommended). + pub fn start_background_sync( + self: &Arc, + base_url: String, + interval: Duration, + ) -> tokio::task::JoinHandle<()> { + let this = Arc::clone(self); tokio::spawn(async move { + // Immediate first sync so the catalog is warm before first dispatch. + if let Err(e) = this.fetch_live(&base_url).await { + tracing::warn!("Initial models registry sync failed: {}", e); + } let mut ticker = tokio::time::interval(interval); + // Skip the immediate first tick: we just synced above. + ticker.tick().await; loop { ticker.tick().await; - if let Err(e) = self.fetch_live(&base_url).await { + if let Err(e) = this.fetch_live(&base_url).await { tracing::warn!("Periodic models registry sync failed: {}", e); } else { tracing::debug!("Periodic models registry sync succeeded"); } } - }); + }) } pub async fn is_model_available(&self, model_id: &str) -> bool { diff --git a/src/router/mod.rs b/src/router/mod.rs index 959c396..3c8c0f4 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -23,30 +23,37 @@ pub struct RouteDecision { pub struct Router { pub default_persona: String, - router_model: std::sync::RwLock, - pub complexity_threshold: u8, + router_model: arc_swap::ArcSwap, + complexity_threshold: std::sync::atomic::AtomicU8, } impl Router { pub fn new(default_persona: String, router_model: String, complexity_threshold: u8) -> Self { Self { default_persona, - router_model: std::sync::RwLock::new(router_model), - complexity_threshold, + router_model: arc_swap::ArcSwap::new(std::sync::Arc::new(router_model)), + complexity_threshold: std::sync::atomic::AtomicU8::new( + complexity_threshold.clamp(1, 10), + ), } } + pub fn complexity_threshold(&self) -> u8 { + self.complexity_threshold + .load(std::sync::atomic::Ordering::Relaxed) + } + pub fn set_router_model(&self, model: String) { - if let Ok(mut write) = self.router_model.write() { - *write = model; - } + self.router_model.store(std::sync::Arc::new(model)); } pub fn get_router_model(&self) -> String { - self.router_model - .read() - .map(|m| m.clone()) - .unwrap_or_default() + self.router_model.load_full().as_ref().clone() + } + + pub fn set_complexity_threshold(&self, threshold: u8) { + self.complexity_threshold + .store(threshold.clamp(1, 10), std::sync::atomic::Ordering::Relaxed); } pub async fn route( @@ -57,7 +64,7 @@ impl Router { force_escalate: bool, ) -> Result { let complexity_score = LlmRouter::estimate_complexity(prompt); - let escalate_to_paid = force_escalate || complexity_score >= self.complexity_threshold; + let escalate_to_paid = force_escalate || complexity_score >= self.complexity_threshold(); // Stage A: Free, instant keyword matching if let Some(matched) = KeywordRouter::classify(registry.list(), prompt) { @@ -156,4 +163,16 @@ mod tests { assert_eq!(dec2.persona, "coder"); assert!(dec2.escalate_to_paid); } + + #[test] + fn test_set_complexity_threshold_clamps_and_reads() { + let router = Router::new("general".into(), "router-model".into(), 7); + assert_eq!(router.complexity_threshold(), 7); + router.set_complexity_threshold(9); + assert_eq!(router.complexity_threshold(), 9); + router.set_complexity_threshold(0); + assert_eq!(router.complexity_threshold(), 1); + router.set_complexity_threshold(255); + assert_eq!(router.complexity_threshold(), 10); + } } diff --git a/src/session.rs b/src/session.rs index 3751d9d..cca4969 100644 --- a/src/session.rs +++ b/src/session.rs @@ -34,8 +34,11 @@ pub struct SessionSummary { impl Session { pub fn new(persona: &str) -> Self { let now: DateTime = Utc::now(); - let timestamp_str = now.format("%Y%m%d-%H%M%S").to_string(); - let id = format!("session-{timestamp_str}"); + let timestamp_str = now.format("%Y%m%d-%H%M%S-%f").to_string(); + // Sub-second timestamp plus random suffix: two sessions created in the + // same microsecond still get distinct ids and never overwrite each other. + let rand_suffix: u32 = rand::random::() % 0xFFFF; + let id = format!("session-{timestamp_str}-{rand_suffix:04x}"); Self { id, @@ -166,4 +169,12 @@ mod tests { assert_eq!(summaries.len(), 1); assert_eq!(summaries[0].id, session.id); } + + #[test] + fn test_session_ids_unique_same_instant() { + // Regression test: two sessions created back-to-back must never collide. + let a = Session::new("coder"); + let b = Session::new("coder"); + assert_ne!(a.id, b.id, "session ids must be unique"); + } } diff --git a/src/tool.rs b/src/tool.rs index 88fb483..c49c9fb 100644 --- a/src/tool.rs +++ b/src/tool.rs @@ -139,8 +139,12 @@ fn truncate_output(s: &str, max_bytes: usize) -> String { if s.len() <= max_bytes { s.to_string() } else { - let truncated = &s[..max_bytes]; - format!("{truncated}\n...[output truncated by BrokeCrew]") + // Walk back to a char boundary so multibyte output never panics. + let mut end = max_bytes.min(s.len()); + while end > 0 && !s.is_char_boundary(end) { + end -= 1; + } + format!("{}\n...[output truncated by BrokeCrew]", &s[..end]) } } @@ -164,4 +168,15 @@ mod tests { assert_eq!(res.exit_code, 0); assert!(res.stdout.contains("hello from brokecrew")); } + + #[test] + fn test_truncate_output_respects_char_boundary() { + // 32KB+ of multibyte output must truncate without panicking. + let big = "日本語テスト出力 ".repeat(4000); + let out = truncate_output(&big, 32 * 1024); + assert!(out.ends_with("...[output truncated by BrokeCrew]")); + assert!(out.len() <= 32 * 1024 + 64); + // Short output passes through untouched. + assert_eq!(truncate_output("ok", 1024), "ok"); + } } diff --git a/src/tui/app.rs b/src/tui/app.rs index a8b459d..2a9b26b 100644 --- a/src/tui/app.rs +++ b/src/tui/app.rs @@ -12,13 +12,6 @@ pub enum MessageSender { model: String, tier: Tier, }, - PipelineStep { - step_index: usize, - total_steps: usize, - persona: String, - model: String, - tier: Tier, - }, Tool { command: String, exit_code: i32, @@ -473,14 +466,8 @@ impl AppState { pub fn update_current_message_model(&mut self, model_id: &str) { if let Some(msg) = self.messages.last_mut() { - match &mut msg.sender { - MessageSender::Persona { model, .. } => { - *model = model_id.to_string(); - } - MessageSender::PipelineStep { model, .. } => { - *model = model_id.to_string(); - } - _ => {} + if let MessageSender::Persona { model, .. } = &mut msg.sender { + *model = model_id.to_string(); } } } @@ -512,30 +499,6 @@ impl AppState { } } - pub fn start_pipeline_step_message( - &mut self, - step_index: usize, - total_steps: usize, - persona: &str, - model: &str, - tier: Tier, - ) { - self.messages.push(ChatMessage { - sender: MessageSender::PipelineStep { - step_index, - total_steps, - persona: persona.to_string(), - model: model.to_string(), - tier, - }, - content: String::new(), - fell_back: false, - }); - if self.auto_scroll { - self.scroll_offset = 0; - } - } - pub fn add_tool_result_message( &mut self, command: String, @@ -586,13 +549,11 @@ impl AppState { let role = match &msg.sender { MessageSender::User => "you", MessageSender::Persona { name, .. } => name.as_str(), - MessageSender::PipelineStep { persona, .. } => persona.as_str(), MessageSender::Tool { .. } => "tool", MessageSender::System => "system", }; let model = match &msg.sender { MessageSender::Persona { model, .. } => Some(model.as_str()), - MessageSender::PipelineStep { model, .. } => Some(model.as_str()), _ => None, }; self.active_session.add_message(role, &msg.content, model); @@ -695,22 +656,6 @@ impl AppState { msg.content.trim() )); } - MessageSender::PipelineStep { - step_index, - total_steps, - persona, - model, - .. - } => { - content.push_str(&format!( - "### Pipeline Step {}/{}: {} ({})\n\n{}\n\n", - step_index + 1, - total_steps, - persona, - model, - msg.content.trim() - )); - } MessageSender::Tool { command, exit_code, diff --git a/src/tui/mod.rs b/src/tui/mod.rs index 395801b..15dc7f0 100644 --- a/src/tui/mod.rs +++ b/src/tui/mod.rs @@ -6,7 +6,7 @@ use brokecrew::budget::BudgetTracker; use brokecrew::client::OpenRouterClient; use brokecrew::config::AppConfig; use brokecrew::persona::Tier; -use brokecrew::pipeline::{PipelineEvent, PipelineExecutor}; +use brokecrew::pipeline::{DispatchParams, PipelineEvent, PipelineExecutor, PipelineRunParams}; use brokecrew::registry::ModelRegistry; use brokecrew::router::{RouteStage, Router}; use brokecrew::session::Session; @@ -66,8 +66,6 @@ pub enum WorkerEvent { step_index: usize, total_steps: usize, persona: String, - model: String, - tier: Tier, }, Chunk { persona: String, @@ -77,6 +75,11 @@ pub enum WorkerEvent { }, StepCompleted { fell_back: bool, + model_used: String, + tier_used: Tier, + }, + StepFailed { + error: String, }, Finished { fell_back: bool, @@ -137,12 +140,11 @@ pub async fn run_tui( app.allow_stealth = allow_stealth || config.allow_stealth; app.escalate = escalate; - // Refresh model catalog in background - let reg_clone = Arc::clone(&model_registry); - let url_clone = config.openrouter_base_url.clone(); - tokio::spawn(async move { - let _ = reg_clone.fetch_live(&url_clone).await; - }); + // Refresh model catalog in background, then keep it warm every 30 min. + let _registry_sync = model_registry.start_background_sync( + config.openrouter_base_url.clone(), + Duration::from_secs(30 * 60), + ); let client = Arc::new(OpenRouterClient::with_base_url( config.api_key.clone().unwrap_or_default(), @@ -313,9 +315,9 @@ fn handle_worker_event(app: &mut AppState, budget_tracker: &mut BudgetTracker, e step_index, total_steps, persona, - model, - tier, } => { + // The model is unknown until the fallback walk resolves; show a + // pending label that CandidateAttempt / StepCompleted will replace. if step_index < app.workings.pipeline_steps.len() { for (i, st) in app.workings.pipeline_steps.iter_mut().enumerate() { if i < step_index { @@ -323,8 +325,9 @@ fn handle_worker_event(app: &mut AppState, budget_tracker: &mut BudgetTracker, e st.is_active = false; } else if i == step_index { st.is_active = true; - st.model = model.clone(); - st.tier = tier; + if st.model.is_empty() { + st.model = "resolving...".to_string(); + } } else { st.is_active = false; } @@ -334,26 +337,21 @@ fn handle_worker_event(app: &mut AppState, budget_tracker: &mut BudgetTracker, e .pipeline_steps .push(app::PipelineStepTelemetry { persona: persona.clone(), - model: model.clone(), - tier, + model: "resolving...".to_string(), + tier: Tier::Free, is_active: true, is_completed: false, fell_back: false, }); } - - if total_steps > 1 { - app.start_pipeline_step_message(step_index, total_steps, &persona, &model, tier); - app.log_working_event(&format!( - "[pipeline] step {}/{}: {} ({})", - step_index + 1, - total_steps, - persona, - model - )); - } else { - app.log_working_event(&format!("[dispatch] {} ({})", persona, model)); - } + app.log_working_event(&format!( + "[pipeline] step {}/{}: {} started", + step_index + 1, + total_steps, + persona + )); + // The chat card opens on first streamed chunk (Chunk handler), + // which carries the real persona/model/tier. } WorkerEvent::Chunk { @@ -376,15 +374,33 @@ fn handle_worker_event(app: &mut AppState, budget_tracker: &mut BudgetTracker, e app.append_to_current_message(&delta); } - WorkerEvent::StepCompleted { fell_back } => { + WorkerEvent::StepCompleted { + fell_back, + model_used, + tier_used, + } => { if let Some(st) = app.workings.pipeline_steps.iter_mut().find(|s| s.is_active) { st.is_completed = true; st.is_active = false; st.fell_back = fell_back; + st.model = model_used.clone(); + st.tier = tier_used; } + app.update_current_message_model(&model_used); if fell_back { app.mark_current_message_fallback(); } + app.log_working_event(&format!("[pipeline] step served by {model_used}")); + } + + WorkerEvent::StepFailed { error } => { + if let Some(st) = app.workings.pipeline_steps.iter_mut().find(|s| s.is_active) { + st.is_completed = true; + st.is_active = false; + st.fell_back = false; + } + app.add_system_message(format!("! pipeline step failed: {error}")); + app.log_working_event(&format!("[pipeline] step failed: {error}")); } WorkerEvent::Finished { @@ -434,10 +450,10 @@ fn handle_worker_event(app: &mut AppState, budget_tracker: &mut BudgetTracker, e let free_calls = calls_made.saturating_sub(paid_calls_made); for _ in 0..free_calls { - let _ = budget_tracker.record_call(false); + let _ = budget_tracker.try_acquire(false); } for _ in 0..paid_calls_made { - let _ = budget_tracker.record_call(true); + let _ = budget_tracker.try_acquire(true); } // Auto-save session state @@ -446,6 +462,16 @@ fn handle_worker_event(app: &mut AppState, budget_tracker: &mut BudgetTracker, e } } +/// Applies a complexity threshold change to in-memory state, the live router, +/// and persisted config in one call, so the /settings slider can never drift. +fn settings_apply_threshold(app: &mut AppState, router: &Arc, next: u8) { + let clamped = next.clamp(1, 10); + app.complexity_threshold = clamped; + router.set_complexity_threshold(clamped); + let _ = AppConfig::save_complexity_threshold(&app.config_dir, clamped); + app.log_working_event(&format!("[settings] complexity threshold: {clamped}")); +} + async fn handle_key_event( app: &mut AppState, key: KeyEvent, @@ -596,15 +622,12 @@ async fn handle_key_event( )); } 3 => { - app.complexity_threshold = if app.complexity_threshold >= 10 { + let next = if app.complexity_threshold >= 10 { 1 } else { app.complexity_threshold + 1 }; - app.log_working_event(&format!( - "[settings] complexity threshold: {}", - app.complexity_threshold - )); + settings_apply_threshold(app, router, next); } 4 => { app.tool_execution_enabled = !app.tool_execution_enabled; @@ -645,11 +668,7 @@ async fn handle_key_event( app.log_working_event(&format!("[settings] theme switched to {theme_name}")); } 3 if app.complexity_threshold > 1 => { - app.complexity_threshold -= 1; - app.log_working_event(&format!( - "[settings] complexity threshold: {}", - app.complexity_threshold - )); + settings_apply_threshold(app, router, app.complexity_threshold - 1); } _ => {} }, @@ -661,11 +680,7 @@ async fn handle_key_event( app.log_working_event(&format!("[settings] theme switched to {theme_name}")); } 3 if app.complexity_threshold < 10 => { - app.complexity_threshold += 1; - app.log_working_event(&format!( - "[settings] complexity threshold: {}", - app.complexity_threshold - )); + settings_apply_threshold(app, router, app.complexity_threshold + 1); } _ => {} }, @@ -1207,19 +1222,18 @@ fn spawn_dispatch_task( step_index, total_steps, persona, - model, - tier, } => { let _ = w_tx .send(WorkerEvent::StepStarted { step_index, total_steps, persona, - model, - tier, }) .await; } + PipelineEvent::StepFailed { error, .. } => { + let _ = w_tx.send(WorkerEvent::StepFailed { error }).await; + } PipelineEvent::CandidateAttempt { model, attempt, @@ -1251,29 +1265,41 @@ fn spawn_dispatch_task( }) .await; } - PipelineEvent::StepCompleted { fell_back, .. } => { - let _ = w_tx.send(WorkerEvent::StepCompleted { fell_back }).await; + PipelineEvent::StepCompleted { + fell_back, + model_used, + tier_used, + .. + } => { + let _ = w_tx + .send(WorkerEvent::StepCompleted { + fell_back, + model_used, + tier_used, + }) + .await; } } } }); - let res = PipelineExecutor::execute( - pipe, - &persona_registry, - &model_registry, - client.as_ref(), - &prompt, + let outcome = PipelineExecutor::execute_outcome(PipelineRunParams { + pipeline: pipe, + persona_registry: &persona_registry, + model_registry: &model_registry, + client: client.as_ref(), + initial_prompt: &prompt, allow_stealth, - escalate, - Some(event_tx), - ) + require_paid: escalate, + event_sender: Some(event_tx), + }) .await; let _ = forwarder.await; - match res { - Ok(results) => { + match outcome.error { + None => { + let results = outcome.results; let total_calls: usize = results.iter().map(|r| r.attempts.len()).sum(); let paid_calls: usize = results.iter().filter(|r| r.tier_used == Tier::Paid).count(); @@ -1295,12 +1321,16 @@ fn spawn_dispatch_task( }) .await; } - Err(e) => { + Some(e) => { + let partial = outcome.results.len(); + let failed_step = outcome.failed_step.unwrap_or(0); let _ = worker_tx .send(WorkerEvent::Finished { fell_back: false, model_used: "none".to_string(), - error: Some(e.to_string()), + error: Some(format!( + "pipeline failed at step {failed_step} after {partial} completed step(s): {e}" + )), calls_made: 1, paid_calls_made: 0, tokens_in: brokecrew::estimate_tokens(&prompt), @@ -1413,14 +1443,16 @@ fn spawn_dispatch_task( let persona_registry_tools = persona_registry.with_tool_capabilities(); let res = PipelineExecutor::dispatch_persona_with_events( - &target_persona, - &persona_registry_tools, - &model_registry, - client.as_ref(), - &prompt, - allow_stealth, - escalate, - Some(chunk_tx), + DispatchParams { + persona_name: &target_persona, + persona_registry: &persona_registry_tools, + model_registry: &model_registry, + client: client.as_ref(), + prompt: &prompt, + allow_stealth, + require_paid: escalate, + chunk_sender: Some(chunk_tx), + }, Some(event_tx), ) .await; @@ -1486,14 +1518,16 @@ fn spawn_dispatch_task( }); let followup_res = PipelineExecutor::dispatch_persona_with_events( - &target_persona, - &persona_registry_tools, - &model_registry, - client.as_ref(), - &followup_prompt, - allow_stealth, - escalate, - Some(chunk_tx2), + DispatchParams { + persona_name: &target_persona, + persona_registry: &persona_registry_tools, + model_registry: &model_registry, + client: client.as_ref(), + prompt: &followup_prompt, + allow_stealth, + require_paid: escalate, + chunk_sender: Some(chunk_tx2), + }, None, ) .await; diff --git a/src/tui/ui.rs b/src/tui/ui.rs index 2bcb0fe..ca4c477 100644 --- a/src/tui/ui.rs +++ b/src/tui/ui.rs @@ -626,43 +626,6 @@ fn render_minimal_messages(frame: &mut Frame, app: &AppState, area: Rect, palett lines.push(Line::from("")); } - MessageSender::PipelineStep { - step_index, - total_steps, - persona, - model, - tier, - } => { - let tier_str = match tier { - Tier::Free => "free", - Tier::Paid => "paid", - }; - let persona_color = get_persona_color(persona); - let mut header_spans = vec![ - Span::raw(" "), - Span::styled( - format!("● Step {}/{}: {persona}", step_index + 1, total_steps), - Style::default().fg(persona_color).bold(), - ), - Span::styled( - format!(" ({model} • {tier_str})"), - Style::default().fg(palette.text_muted), - ), - ]; - if msg.fell_back { - header_spans.push(Span::raw(" ")); - header_spans.push(Span::styled( - "[fallback]", - Style::default().fg(Color::Yellow).bold(), - )); - } - lines.push(Line::from(header_spans)); - lines.push(Line::from("")); - - render_clean_markdown_body(&msg.content, &mut lines); - lines.push(Line::from("")); - } - MessageSender::Tool { command, exit_code, diff --git a/tests/fallback_tests.rs b/tests/fallback_tests.rs index 891dd4b..18160cb 100644 --- a/tests/fallback_tests.rs +++ b/tests/fallback_tests.rs @@ -1,7 +1,7 @@ use async_trait::async_trait; use brokecrew::client::{ClientError, ModelClient}; use brokecrew::persona::{ModelCandidate, Persona, PersonaRegistry, Tier}; -use brokecrew::pipeline::PipelineExecutor; +use brokecrew::pipeline::{DispatchParams, PipelineExecutor}; use brokecrew::registry::ModelRegistry; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; @@ -84,16 +84,16 @@ async fn test_candidate_cascade_on_503_service_unavailable() { calls: calls.clone(), }; - let result = PipelineExecutor::dispatch_persona( - "resilient-coder", - &persona_reg, - &model_reg, - &client, - "write rust code", - false, - false, - None, - ) + let result = PipelineExecutor::dispatch_persona(DispatchParams { + persona_name: "resilient-coder", + persona_registry: &persona_reg, + model_registry: &model_reg, + client: &client, + prompt: "write rust code", + allow_stealth: false, + require_paid: false, + chunk_sender: None, + }) .await .expect("fallback to candidate-2 should succeed"); @@ -121,16 +121,16 @@ async fn test_candidate_cascade_on_429_rate_limited() { calls: calls.clone(), }; - let result = PipelineExecutor::dispatch_persona( - "resilient-coder", - &persona_reg, - &model_reg, - &client, - "write python code", - false, - false, - None, - ) + let result = PipelineExecutor::dispatch_persona(DispatchParams { + persona_name: "resilient-coder", + persona_registry: &persona_reg, + model_registry: &model_reg, + client: &client, + prompt: "write python code", + allow_stealth: false, + require_paid: false, + chunk_sender: None, + }) .await .expect("fallback to candidate-2 should succeed"); @@ -151,16 +151,16 @@ async fn test_fail_fast_on_401_unauthorized() { calls: calls.clone(), }; - let err = PipelineExecutor::dispatch_persona( - "resilient-coder", - &persona_reg, - &model_reg, - &client, - "test query", - false, - false, - None, - ) + let err = PipelineExecutor::dispatch_persona(DispatchParams { + persona_name: "resilient-coder", + persona_registry: &persona_reg, + model_registry: &model_reg, + client: &client, + prompt: "test query", + allow_stealth: false, + require_paid: false, + chunk_sender: None, + }) .await .unwrap_err(); @@ -208,16 +208,16 @@ async fn test_total_cascade_failure_reports_all_attempts() { } } - let err = PipelineExecutor::dispatch_persona( - "resilient-coder", - &persona_reg, - &model_reg, - &AllFailingClient, - "test query", - false, - false, - None, - ) + let err = PipelineExecutor::dispatch_persona(DispatchParams { + persona_name: "resilient-coder", + persona_registry: &persona_reg, + model_registry: &model_reg, + client: &AllFailingClient, + prompt: "test query", + allow_stealth: false, + require_paid: false, + chunk_sender: None, + }) .await .unwrap_err();