mod cli; use brokecrew::budget::BudgetTracker; use brokecrew::client::OpenRouterClient; use brokecrew::config::AppConfig; use brokecrew::persona::Tier; use brokecrew::pipeline::{PipelineConfig, PipelineEvent, PipelineExecutor}; use brokecrew::registry::ModelRegistry; use brokecrew::router::{RouteStage, Router}; use clap::Parser; use cli::{ChatArgs, Cli, Commands, ConfigSubcommand, PersonasSubcommand}; use indicatif::{ProgressBar, ProgressStyle}; use owo_colors::OwoColorize; use std::io::{self, BufRead, Write}; use std::sync::Arc; use std::time::Duration; use tokio::sync::mpsc; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter}; fn init_logging(json_logs: bool, verbose: bool) { let filter = if verbose { EnvFilter::new("brokecrew=debug,reqwest=debug") } else { EnvFilter::new("brokecrew=info") }; if json_logs { tracing_subscriber::registry() .with(filter) .with(tracing_subscriber::fmt::layer().json()) .init(); } else { tracing_subscriber::registry() .with(filter) .with( tracing_subscriber::fmt::layer() .with_writer(io::stderr) .without_time() .with_target(false), ) .init(); } } fn get_persona_color_tag(persona: &str) -> String { match persona.to_lowercase().as_str() { "coder" => format!("[{}]", persona).bright_green().bold().to_string(), "reviewer" => format!("[{}]", persona).bright_yellow().bold().to_string(), "architect" => format!("[{}]", persona).bright_blue().bold().to_string(), "general" | "assistant" => format!("[{}]", persona).bright_magenta().bold().to_string(), _ => format!("[{}]", persona).bright_cyan().bold().to_string(), } } fn create_spinner(message: &'static str) -> ProgressBar { let pb = ProgressBar::new_spinner(); pb.set_style( ProgressStyle::default_spinner() .tick_chars("⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏ ") .template("{spinner:.cyan} {msg}") .expect("valid template"), ); pb.set_message(message); pb.enable_steady_tick(Duration::from_millis(80)); pb } #[tokio::main] async fn main() -> Result<(), Box> { let cli = Cli::parse(); init_logging(cli.json_logs, cli.verbose); let config = AppConfig::load_from_paths( cli.config_file.as_deref(), cli.personas_file.as_deref(), cli.api_key.clone(), ) .map_err(|e| format!("Configuration error: {e}"))?; let budget_path = config.config_dir.join("budget.json"); let mut budget_tracker = BudgetTracker::new(Some(budget_path), config.credit_tier); let model_registry = Arc::new(ModelRegistry::new()); // Check subcommand match cli.command { Some(Commands::Config(config_args)) => match config_args.action { ConfigSubcommand::SetKey { key } => { let api_key = match key { Some(k) if !k.trim().is_empty() => k.trim().to_string(), _ => { eprint!("Enter OpenRouter API Key: "); io::stderr().flush()?; let mut input = String::new(); io::stdin().read_line(&mut input)?; let trimmed = input.trim().to_string(); if trimmed.is_empty() { eprintln!("{}", "Error: API key cannot be empty.".red().bold()); std::process::exit(1); } trimmed } }; 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() ); return Ok(()); } }, Some(Commands::Personas(personas_args)) => match personas_args.action { PersonasSubcommand::List => { println!("\n{}", "BrokeCrew Personas:".bold().underline()); for persona in config.persona_registry.list() { 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!(" Candidates:"); for (i, c) in persona.candidates.iter().enumerate() { let tier_tag = match c.tier { Tier::Free => "FREE".bright_green().to_string(), Tier::Paid => "PAID".bright_yellow().to_string(), }; let stealth_tag = if c.is_stealth { format!(" {}", "[stealth]".bright_magenta()) } else { String::new() }; println!( " {}. {} ({}){}", i + 1, c.id.bold(), tier_tag, stealth_tag ); } } if !config.pipelines.is_empty() { println!("\n{}", "Configured Pipelines:".bold().underline()); for pipe in &config.pipelines { let steps_str: Vec = pipe .steps .iter() .map(|s| { if s.pass_previous_output { format!("{} (with context)", s.persona.bold()) } else { s.persona.bold().to_string() } }) .collect(); println!( " • {} ➔ {}", pipe.name.bright_cyan().bold(), steps_str.join(" ➔ ") ); } } println!(); return Ok(()); } }, Some(Commands::Status) => { println!("\n{}", "BrokeCrew Status & Rate Limits:".bold().underline()); // Fetch live registry for status count let spinner = create_spinner("Fetching live OpenRouter model catalog..."); let model_count = match model_registry.fetch_live(&config.openrouter_base_url).await { Ok(count) => { spinner.finish_and_clear(); count } Err(e) => { spinner.finish_and_clear(); println!(" Live catalog: {}", format!("Offline / Error: {e}").red()); 0 } }; if model_count > 0 { 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() } else { 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!(" 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()); } } Err(e) => { println!(" Budget status: {}", e.to_string().red().bold()); } } println!(); return Ok(()); } Some(Commands::Chat(args)) => { run_chat(args, &config, &mut budget_tracker, model_registry).await?; } None => { // Default to chat if no subcommand passed run_chat( ChatArgs::default(), &config, &mut budget_tracker, model_registry, ) .await?; } } Ok(()) } async fn run_chat( args: ChatArgs, config: &AppConfig, budget_tracker: &mut BudgetTracker, model_registry: Arc, ) -> Result<(), Box> { let api_key = match config.api_key.clone() { Some(k) if !k.trim().is_empty() => k, _ => { eprintln!( "{}\nRun: {}\nOr set: {}", "Error: OpenRouter API key not configured.".red().bold(), "brokecrew config set-key".bright_green().bold(), "export OPENROUTER_API_KEY=\"sk-or-v1-...\"".cyan() ); std::process::exit(1); } }; 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; }); let router = Router::new( config.default_persona.clone(), config.router_model.clone(), config.complexity_threshold, ); // If prompt passed directly as CLI argument, execute single-shot if let Some(ref prompt) = args.prompt { execute_prompt( prompt, &args, config, budget_tracker, &model_registry, &router, &client, ) .await?; return Ok(()); } // 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() ); let stdin = io::stdin(); let mut reader = stdin.lock(); loop { print!("{} ", "you ❯".bright_white().bold()); io::stdout().flush()?; let mut line = String::new(); if reader.read_line(&mut line)? == 0 { break; } let trimmed = line.trim(); if trimmed.is_empty() { continue; } if trimmed.eq_ignore_ascii_case("exit") || trimmed.eq_ignore_ascii_case("quit") { break; } execute_prompt( trimmed, &args, config, budget_tracker, &model_registry, &router, &client, ) .await?; println!(); } Ok(()) } async fn execute_prompt( prompt: &str, args: &ChatArgs, config: &AppConfig, budget_tracker: &mut BudgetTracker, model_registry: &ModelRegistry, router: &Router, client: &OpenRouterClient, ) -> Result<(), Box> { // 1. Budget check match budget_tracker.check_budget() { Ok(status) => { if status.approaching_daily_limit { eprintln!( "{} Daily usage: {}/{}", "⚠ Budget notice: Approaching daily ceiling.".yellow(), status.daily_requests, status.daily_limit ); } } Err(e) => { eprintln!("{} {}", "Rate limit warning:".red().bold(), e); return Ok(()); } } let allow_stealth = args.allow_stealth || config.allow_stealth; let require_paid = args.escalate; // 2. Check if running a multi-agent pipeline if let Some(ref pipe_name) = args.pipeline { let pipeline = config .pipelines .iter() .find(|p| p.name == *pipe_name) .ok_or_else(|| format!("Pipeline '{pipe_name}' not found in configuration"))?; run_pipeline( pipeline, config, budget_tracker, model_registry, client, prompt, allow_stealth, require_paid, ) .await?; return Ok(()); } // 3. Single persona dispatch let (target_persona, route_info) = if let Some(ref p) = args.persona { (p.clone(), None) } else { let spinner = create_spinner("Routing message..."); // Record budget call if LLM router is invoked let decision = router .route(&config.persona_registry, client, prompt, require_paid) .await; spinner.finish_and_clear(); match decision { Ok(dec) => { if dec.stage_used == RouteStage::StageBLlm { let _ = budget_tracker.record_call(false); } let info = match dec.stage_used { RouteStage::StageAKeyword => "routed via keyword pre-filter".to_string(), RouteStage::StageBLlm => "routed via LLM classifier".to_string(), }; (dec.persona, Some(info)) } Err(e) => { 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() ); } let (chunk_tx, mut chunk_rx) = mpsc::channel::(100); 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, model_registry, client, prompt, allow_stealth, require_paid, Some(chunk_tx), ); let printer = tokio::spawn(async move { let mut first = true; while let Some(chunk) = chunk_rx.recv().await { if first { spinner.finish_and_clear(); println!("\n{}", persona_tag); first = false; } print!("{chunk}"); let _ = io::stdout().flush(); } }); let dispatch_result = dispatch_future.await; let _ = printer.await; match dispatch_result { Ok(res) => { println!(); if res.fell_back { eprintln!( "{}", 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) => { eprintln!("\n{} {}", "Error dispatching to persona:".red().bold(), e); } } 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> { println!( "\n{} Starting pipeline '{}' ({} steps)", "⚡".bright_cyan(), pipeline.name.bold(), pipeline.steps.len() ); let (event_tx, mut event_rx) = mpsc::channel::(100); let exec_future = PipelineExecutor::execute( pipeline, &config.persona_registry, model_registry, client, prompt, allow_stealth, require_paid, Some(event_tx), ); let event_handler = tokio::spawn(async move { let mut current_spinner: Option = None; while let Some(event) = event_rx.recv().await { match event { PipelineEvent::StepStarted { 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 {}/{} {} ({}) ---", step_index + 1, total_steps, badge, model.dimmed() ); let sp = create_spinner("Thinking..."); current_spinner = Some(sp); } PipelineEvent::Chunk { delta, .. } => { if let Some(sp) = current_spinner.take() { sp.finish_and_clear(); } print!("{delta}"); let _ = io::stdout().flush(); } PipelineEvent::StepCompleted { fell_back, .. } => { if let Some(sp) = current_spinner.take() { sp.finish_and_clear(); } println!(); if fell_back { eprintln!("{}", "⚠ Fallback occurred during step execution".yellow()); } } } } }); let results = exec_future.await?; 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); } Ok(()) }