Compare commits
15 commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 9fb0874b61 | |||
| 39201f5553 | |||
| 4c0802b172 | |||
| b2f4485946 | |||
| 276e57b2ba | |||
| b1aac41a56 | |||
| e41a641c94 | |||
| 659de0dc04 | |||
| e8dd01509c | |||
| bf02bf3bc6 | |||
| 147d3fb481 | |||
| 7d6fd2d84f | |||
| eb4399177e | |||
| 9e96eec964 | |||
| bfea14d01b |
21 changed files with 4686 additions and 700 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -2,3 +2,4 @@
|
||||||
.superpowers/
|
.superpowers/
|
||||||
**/*.rs.bk
|
**/*.rs.bk
|
||||||
Cargo.lock
|
Cargo.lock
|
||||||
|
PROJECT_CONTEXT.md
|
||||||
|
|
|
||||||
68
.gitlab-ci.yml
Normal file
68
.gitlab-ci.yml
Normal file
|
|
@ -0,0 +1,68 @@
|
||||||
|
stages:
|
||||||
|
- check
|
||||||
|
- build
|
||||||
|
|
||||||
|
variables:
|
||||||
|
CARGO_TERM_COLOR: always
|
||||||
|
RUST_BACKTRACE: "1"
|
||||||
|
|
||||||
|
cache:
|
||||||
|
key:
|
||||||
|
files:
|
||||||
|
- Cargo.lock
|
||||||
|
prefix: cargo
|
||||||
|
paths:
|
||||||
|
- target/
|
||||||
|
|
||||||
|
fmt:
|
||||||
|
stage: check
|
||||||
|
image: rust:latest
|
||||||
|
script:
|
||||||
|
- rustup component add rustfmt
|
||||||
|
- cargo fmt --check
|
||||||
|
rules:
|
||||||
|
- if: $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH
|
||||||
|
|
||||||
|
clippy:
|
||||||
|
stage: check
|
||||||
|
image: rust:latest
|
||||||
|
before_script:
|
||||||
|
- rustup component add clippy
|
||||||
|
script:
|
||||||
|
- cargo clippy --all-targets -- --deny warnings
|
||||||
|
rules:
|
||||||
|
- if: $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH
|
||||||
|
allow_failure: true
|
||||||
|
|
||||||
|
test:
|
||||||
|
stage: check
|
||||||
|
image: rust:latest
|
||||||
|
script:
|
||||||
|
- cargo test --locked
|
||||||
|
rules:
|
||||||
|
- if: $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH
|
||||||
|
|
||||||
|
build:
|
||||||
|
stage: build
|
||||||
|
image: rust:latest
|
||||||
|
script:
|
||||||
|
- cargo build --release --locked
|
||||||
|
artifacts:
|
||||||
|
name: "brokecream-linux-x86_64"
|
||||||
|
paths:
|
||||||
|
- target/release/brokecrew
|
||||||
|
expire_in: 30 days
|
||||||
|
|
||||||
|
# Only on tags, so routine pushes don't eat the free CI minutes.
|
||||||
|
release:
|
||||||
|
stage: build
|
||||||
|
image: rust:latest
|
||||||
|
rules:
|
||||||
|
- if: $CI_COMMIT_TAG
|
||||||
|
script:
|
||||||
|
- cargo build --release --locked
|
||||||
|
artifacts:
|
||||||
|
name: "brokecrew-$CI_COMMIT_TAG"
|
||||||
|
paths:
|
||||||
|
- target/release/brokecrew
|
||||||
|
expire_in: 90 days
|
||||||
|
|
@ -28,6 +28,8 @@ owo-colors = "4.1"
|
||||||
regex = "1.10"
|
regex = "1.10"
|
||||||
chrono = { version = "0.4", features = ["serde"] }
|
chrono = { version = "0.4", features = ["serde"] }
|
||||||
dirs = "5.0"
|
dirs = "5.0"
|
||||||
|
rand = "0.8"
|
||||||
|
arc-swap = "1.7"
|
||||||
tempfile = "3.12"
|
tempfile = "3.12"
|
||||||
ratatui = "0.29"
|
ratatui = "0.29"
|
||||||
crossterm = { version = "0.28", features = ["event-stream"] }
|
crossterm = { version = "0.28", features = ["event-stream"] }
|
||||||
|
|
|
||||||
16
README.md
16
README.md
|
|
@ -1,6 +1,5 @@
|
||||||
# BrokeCrew 🤖⚡
|
# BrokeCrew
|
||||||
|
|
||||||
[](https://github.com/oiupoyt/brokecrew/actions)
|
|
||||||
[](LICENSE-MIT)
|
[](LICENSE-MIT)
|
||||||
[](https://crates.io/crates/brokecrew)
|
[](https://crates.io/crates/brokecrew)
|
||||||
|
|
||||||
|
|
@ -23,7 +22,14 @@ BrokeCrew makes free-tier models first-class citizens while eliminating their in
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## OpenCode / Pi-Style Interactive TUI
|
### Technical Invariants
|
||||||
|
- **Fail-Fast Cascade Walk**: Free endpoints on OpenRouter frequently return transient `429`, `503`, or drop out of `/v1/models` without warning. Candidate resolution treats failures as cascade triggers rather than fatal aborts, stepping to the next viable candidate model within `<250ms`.
|
||||||
|
- **Sliding-Window Budgeting**: Token usage and request counts are stored in atomic sliding windows (`VecDeque<DateTime<Utc>>`). Router and classifier calls deduct from the shared account budget to prevent downstream API lockouts.
|
||||||
|
- **Async Non-Blocking SSE**: Streaming tokens are decoupled from the Ratatui TUI render loop via crossbeam/tokio MPSC channels, ensuring UI frame rates remain at 60 FPS even during erratic remote network jitter.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Brokecrew Interactive TUI
|
||||||
|
|
||||||
Simply run `brokecrew` to enter the rich, full-screen interactive TUI:
|
Simply run `brokecrew` to enter the rich, full-screen interactive TUI:
|
||||||
|
|
||||||
|
|
@ -31,8 +37,8 @@ Simply run `brokecrew` to enter the rich, full-screen interactive TUI:
|
||||||
brokecrew
|
brokecrew
|
||||||
```
|
```
|
||||||
|
|
||||||
- **Live Team View**: Color-coded persona avatars (`⚡ coder`, `🔍 reviewer`, `🏛️ architect`), token streaming with markdown & syntax formatting.
|
- **Live Team View**: Color-coded persona indicators (`● coder`, `● reviewer`, `● architect`), token streaming with markdown & syntax formatting.
|
||||||
- **Quota & Budget Monitor**: Real-time headers displaying rolling minute and 24h limits (`⏱ 1m: 2/20`, `📅 24h: 14/50`).
|
- **Quota & Budget Monitor**: Real-time headers displaying rolling minute and 24h limits (`1m: 2/20`, `24h: 14/50`).
|
||||||
- **Interactive Modals**:
|
- **Interactive Modals**:
|
||||||
- `[Tab]` — Open Persona Selector (switch between intelligent Stage A/B `Auto-Route` and dedicated personas).
|
- `[Tab]` — Open Persona Selector (switch between intelligent Stage A/B `Auto-Route` and dedicated personas).
|
||||||
- `[Ctrl+P]` — Open Pipeline Selector (execute multi-agent pipelines like `build-and-review`).
|
- `[Ctrl+P]` — Open Pipeline Selector (execute multi-agent pipelines like `build-and-review`).
|
||||||
|
|
|
||||||
|
|
@ -77,10 +77,10 @@ is_stealth = false
|
||||||
[[persona]]
|
[[persona]]
|
||||||
name = "general"
|
name = "general"
|
||||||
system_prompt = "You are BrokeCrew's general assistant: insightful, direct, and concise."
|
system_prompt = "You are BrokeCrew's general assistant: insightful, direct, and concise."
|
||||||
trigger_words = ["help", "explain", "summarize", "what", "who", "why", "how", "hello"]
|
trigger_words = ["help", "explain", "summarize", "what", "who", "why", "how", "hello", "hi", "hey", "yo", "sup", "greetings"]
|
||||||
|
|
||||||
[[persona.candidates]]
|
[[persona.candidates]]
|
||||||
id = "qwen/qwen3.8-27b:free"
|
id = "liquid/lfm-2.5-2.6b:free"
|
||||||
tier = "free"
|
tier = "free"
|
||||||
is_stealth = false
|
is_stealth = false
|
||||||
|
|
||||||
|
|
@ -90,7 +90,7 @@ tier = "free"
|
||||||
is_stealth = false
|
is_stealth = false
|
||||||
|
|
||||||
[[persona.candidates]]
|
[[persona.candidates]]
|
||||||
id = "liquid/lfm-2.5-2.6b:free"
|
id = "qwen/qwen3.8-27b:free"
|
||||||
tier = "free"
|
tier = "free"
|
||||||
is_stealth = false
|
is_stealth = false
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -7,6 +7,14 @@ use std::path::{Path, PathBuf};
|
||||||
|
|
||||||
const MINUTE_LIMIT: usize = 20;
|
const MINUTE_LIMIT: usize = 20;
|
||||||
|
|
||||||
|
/// Estimates token count from text using standard heuristic (~4 characters per token).
|
||||||
|
pub fn estimate_tokens(text: &str) -> usize {
|
||||||
|
if text.is_empty() {
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
|
text.len().div_ceil(4).max(1)
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
pub enum CreditTier {
|
pub enum CreditTier {
|
||||||
LowBalance, // < $10 lifetime purchased: 50 requests/day
|
LowBalance, // < $10 lifetime purchased: 50 requests/day
|
||||||
|
|
@ -28,6 +36,7 @@ pub struct BudgetStatus {
|
||||||
pub minute_limit: usize,
|
pub minute_limit: usize,
|
||||||
pub daily_requests: usize,
|
pub daily_requests: usize,
|
||||||
pub daily_limit: usize,
|
pub daily_limit: usize,
|
||||||
|
pub daily_paid_requests: usize,
|
||||||
pub approaching_daily_limit: bool,
|
pub approaching_daily_limit: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -42,6 +51,7 @@ pub enum BudgetExceeded {
|
||||||
#[derive(Debug, Serialize, Deserialize, Default)]
|
#[derive(Debug, Serialize, Deserialize, Default)]
|
||||||
struct StoredBudgetState {
|
struct StoredBudgetState {
|
||||||
daily_count: usize,
|
daily_count: usize,
|
||||||
|
daily_paid_count: usize,
|
||||||
last_day_utc: String, // YYYY-MM-DD
|
last_day_utc: String, // YYYY-MM-DD
|
||||||
recent_calls: VecDeque<DateTime<Utc>>,
|
recent_calls: VecDeque<DateTime<Utc>>,
|
||||||
}
|
}
|
||||||
|
|
@ -50,6 +60,7 @@ pub struct BudgetTracker {
|
||||||
storage_path: Option<PathBuf>,
|
storage_path: Option<PathBuf>,
|
||||||
credit_tier: CreditTier,
|
credit_tier: CreditTier,
|
||||||
daily_count: usize,
|
daily_count: usize,
|
||||||
|
daily_paid_count: usize,
|
||||||
last_day_utc: String,
|
last_day_utc: String,
|
||||||
recent_calls: VecDeque<DateTime<Utc>>,
|
recent_calls: VecDeque<DateTime<Utc>>,
|
||||||
}
|
}
|
||||||
|
|
@ -60,6 +71,7 @@ impl BudgetTracker {
|
||||||
storage_path,
|
storage_path,
|
||||||
credit_tier,
|
credit_tier,
|
||||||
daily_count: 0,
|
daily_count: 0,
|
||||||
|
daily_paid_count: 0,
|
||||||
last_day_utc: Utc::now().format("%Y-%m-%d").to_string(),
|
last_day_utc: Utc::now().format("%Y-%m-%d").to_string(),
|
||||||
recent_calls: VecDeque::new(),
|
recent_calls: VecDeque::new(),
|
||||||
};
|
};
|
||||||
|
|
@ -83,6 +95,7 @@ impl BudgetTracker {
|
||||||
// Day rollover check
|
// Day rollover check
|
||||||
if self.last_day_utc != current_day {
|
if self.last_day_utc != current_day {
|
||||||
self.daily_count = 0;
|
self.daily_count = 0;
|
||||||
|
self.daily_paid_count = 0;
|
||||||
self.last_day_utc = current_day;
|
self.last_day_utc = current_day;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -124,16 +137,49 @@ impl BudgetTracker {
|
||||||
minute_limit: MINUTE_LIMIT,
|
minute_limit: MINUTE_LIMIT,
|
||||||
daily_requests: self.daily_count,
|
daily_requests: self.daily_count,
|
||||||
daily_limit,
|
daily_limit,
|
||||||
|
daily_paid_requests: self.daily_paid_count,
|
||||||
approaching_daily_limit,
|
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<BudgetStatus, BudgetExceeded> {
|
||||||
|
// 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.
|
/// 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();
|
self.refresh_window();
|
||||||
let now = Utc::now();
|
let now = Utc::now();
|
||||||
self.recent_calls.push_back(now);
|
self.recent_calls.push_back(now);
|
||||||
self.daily_count += 1;
|
self.daily_count += 1;
|
||||||
|
if is_paid {
|
||||||
|
self.daily_paid_count += 1;
|
||||||
|
}
|
||||||
|
|
||||||
if let Some(ref path) = self.storage_path.clone() {
|
if let Some(ref path) = self.storage_path.clone() {
|
||||||
self.save_to_disk(path)?;
|
self.save_to_disk(path)?;
|
||||||
|
|
@ -148,6 +194,7 @@ impl BudgetTracker {
|
||||||
let data = fs::read_to_string(path)?;
|
let data = fs::read_to_string(path)?;
|
||||||
if let Ok(state) = serde_json::from_str::<StoredBudgetState>(&data) {
|
if let Ok(state) = serde_json::from_str::<StoredBudgetState>(&data) {
|
||||||
self.daily_count = state.daily_count;
|
self.daily_count = state.daily_count;
|
||||||
|
self.daily_paid_count = state.daily_paid_count;
|
||||||
self.last_day_utc = state.last_day_utc;
|
self.last_day_utc = state.last_day_utc;
|
||||||
self.recent_calls = state.recent_calls;
|
self.recent_calls = state.recent_calls;
|
||||||
}
|
}
|
||||||
|
|
@ -160,6 +207,7 @@ impl BudgetTracker {
|
||||||
}
|
}
|
||||||
let state = StoredBudgetState {
|
let state = StoredBudgetState {
|
||||||
daily_count: self.daily_count,
|
daily_count: self.daily_count,
|
||||||
|
daily_paid_count: self.daily_paid_count,
|
||||||
last_day_utc: self.last_day_utc.clone(),
|
last_day_utc: self.last_day_utc.clone(),
|
||||||
recent_calls: self.recent_calls.clone(),
|
recent_calls: self.recent_calls.clone(),
|
||||||
};
|
};
|
||||||
|
|
@ -237,4 +285,36 @@ mod tests {
|
||||||
assert_eq!(status.daily_requests, 2);
|
assert_eq!(status.daily_requests, 2);
|
||||||
assert_eq!(status.minute_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:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
14
src/cli.rs
14
src/cli.rs
|
|
@ -98,6 +98,20 @@ pub enum ConfigSubcommand {
|
||||||
#[arg(help = "Optional key directly passed; if omitted, securely prompted")]
|
#[arg(help = "Optional key directly passed; if omitted, securely prompted")]
|
||||||
key: Option<String>,
|
key: Option<String>,
|
||||||
},
|
},
|
||||||
|
#[command(about = "Set and persist provider base URL")]
|
||||||
|
SetUrl {
|
||||||
|
#[arg(
|
||||||
|
help = "Provider API base URL (e.g. https://openrouter.ai/api/v1 or http://localhost:11434/v1)"
|
||||||
|
)]
|
||||||
|
url: String,
|
||||||
|
},
|
||||||
|
#[command(about = "Set and persist Stage B router model")]
|
||||||
|
SetRouterModel {
|
||||||
|
#[arg(help = "Model identifier to use for Stage B classification")]
|
||||||
|
model: String,
|
||||||
|
},
|
||||||
|
#[command(about = "Display active configuration and credentials status")]
|
||||||
|
Show,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Args, Debug)]
|
#[derive(Args, Debug)]
|
||||||
|
|
|
||||||
101
src/client.rs
101
src/client.rs
|
|
@ -3,6 +3,7 @@ use futures_util::StreamExt;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
use std::fmt;
|
use std::fmt;
|
||||||
|
use std::sync::Arc;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
use tokio::sync::mpsc;
|
use tokio::sync::mpsc;
|
||||||
|
|
||||||
|
|
@ -37,9 +38,17 @@ pub trait ModelClient: Send + Sync {
|
||||||
) -> Result<String, ClientError>;
|
) -> Result<String, ClientError>;
|
||||||
}
|
}
|
||||||
|
|
||||||
pub struct OpenRouterClient {
|
/// 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,
|
api_key: String,
|
||||||
base_url: String,
|
base_url: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct OpenRouterClient {
|
||||||
|
config: Arc<arc_swap::ArcSwap<ClientConfig>>,
|
||||||
http: reqwest::Client,
|
http: reqwest::Client,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -47,7 +56,7 @@ impl fmt::Debug for OpenRouterClient {
|
||||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||||
f.debug_struct("OpenRouterClient")
|
f.debug_struct("OpenRouterClient")
|
||||||
.field("api_key", &"[REDACTED]")
|
.field("api_key", &"[REDACTED]")
|
||||||
.field("base_url", &self.base_url)
|
.field("base_url", &self.config.load().base_url)
|
||||||
.finish()
|
.finish()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -97,12 +106,38 @@ impl OpenRouterClient {
|
||||||
.unwrap_or_else(|_| reqwest::Client::new());
|
.unwrap_or_else(|_| reqwest::Client::new());
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
api_key,
|
config: Arc::new(arc_swap::ArcSwap::new(Arc::new(ClientConfig {
|
||||||
base_url,
|
api_key,
|
||||||
|
base_url,
|
||||||
|
}))),
|
||||||
http,
|
http,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn set_api_key(&self, key: String) {
|
||||||
|
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) {
|
||||||
|
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.config.load().api_key.clone()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn get_base_url(&self) -> String {
|
||||||
|
self.config.load().base_url.clone()
|
||||||
|
}
|
||||||
|
|
||||||
/// Parses an SSE formatted line into a token delta if available
|
/// Parses an SSE formatted line into a token delta if available
|
||||||
pub fn parse_sse_line(line: &str) -> Option<String> {
|
pub fn parse_sse_line(line: &str) -> Option<String> {
|
||||||
let trimmed = line.trim();
|
let trimmed = line.trim();
|
||||||
|
|
@ -136,31 +171,30 @@ impl OpenRouterClient {
|
||||||
None
|
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(
|
fn check_status(
|
||||||
status: reqwest::StatusCode,
|
status: reqwest::StatusCode,
|
||||||
body: &str,
|
body: &str,
|
||||||
headers: &reqwest::header::HeaderMap,
|
headers: &reqwest::header::HeaderMap,
|
||||||
) -> Result<(), ClientError> {
|
) -> Option<ClientError> {
|
||||||
match status.as_u16() {
|
match status.as_u16() {
|
||||||
200..=299 => Ok(()),
|
200..=299 => None,
|
||||||
401 => Err(ClientError::Unauthorized),
|
401 => Some(ClientError::Unauthorized),
|
||||||
402 => Err(ClientError::PaymentRequired),
|
402 => Some(ClientError::PaymentRequired),
|
||||||
429 => {
|
429 => {
|
||||||
let retry_after = headers
|
let retry_after = headers
|
||||||
.get("retry-after")
|
.get("retry-after")
|
||||||
.and_then(|h| h.to_str().ok())
|
.and_then(|h| h.to_str().ok())
|
||||||
.and_then(|s| s.parse::<u64>().ok())
|
.and_then(|s| s.parse::<u64>().ok())
|
||||||
.unwrap_or(5);
|
.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}"
|
"Model not found (404): {body}"
|
||||||
))),
|
))),
|
||||||
s if s >= 500 => Err(ClientError::ServerError {
|
s => Some(ClientError::ServerError {
|
||||||
status: s,
|
|
||||||
message: body.to_string(),
|
|
||||||
}),
|
|
||||||
s => Err(ClientError::ServerError {
|
|
||||||
status: s,
|
status: s,
|
||||||
message: body.to_string(),
|
message: body.to_string(),
|
||||||
}),
|
}),
|
||||||
|
|
@ -175,7 +209,9 @@ impl OpenRouterClient {
|
||||||
max_tokens: Option<u32>,
|
max_tokens: Option<u32>,
|
||||||
temperature: Option<f32>,
|
temperature: Option<f32>,
|
||||||
) -> Result<String, ClientError> {
|
) -> Result<String, ClientError> {
|
||||||
let endpoint = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
|
let base_url = self.get_base_url();
|
||||||
|
let api_key = self.get_api_key();
|
||||||
|
let endpoint = format!("{}/chat/completions", base_url.trim_end_matches('/'));
|
||||||
|
|
||||||
let body = ChatCompletionRequest {
|
let body = ChatCompletionRequest {
|
||||||
model,
|
model,
|
||||||
|
|
@ -194,9 +230,8 @@ impl OpenRouterClient {
|
||||||
temperature,
|
temperature,
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut auth_val =
|
let mut auth_val = reqwest::header::HeaderValue::from_str(&format!("Bearer {}", api_key))
|
||||||
reqwest::header::HeaderValue::from_str(&format!("Bearer {}", self.api_key))
|
.unwrap_or_else(|_| reqwest::header::HeaderValue::from_static(""));
|
||||||
.unwrap_or_else(|_| reqwest::header::HeaderValue::from_static(""));
|
|
||||||
auth_val.set_sensitive(true);
|
auth_val.set_sensitive(true);
|
||||||
|
|
||||||
let res = self
|
let res = self
|
||||||
|
|
@ -213,9 +248,10 @@ impl OpenRouterClient {
|
||||||
let status = res.status();
|
let status = res.status();
|
||||||
let headers = res.headers().clone();
|
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();
|
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?;
|
let resp_json: ChatCompletionResponse = res.json().await?;
|
||||||
|
|
@ -245,7 +281,9 @@ impl ModelClient for OpenRouterClient {
|
||||||
user: &str,
|
user: &str,
|
||||||
chunk_sender: mpsc::Sender<String>,
|
chunk_sender: mpsc::Sender<String>,
|
||||||
) -> Result<String, ClientError> {
|
) -> Result<String, ClientError> {
|
||||||
let endpoint = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
|
let base_url = self.get_base_url();
|
||||||
|
let api_key = self.get_api_key();
|
||||||
|
let endpoint = format!("{}/chat/completions", base_url.trim_end_matches('/'));
|
||||||
|
|
||||||
let body = ChatCompletionRequest {
|
let body = ChatCompletionRequest {
|
||||||
model,
|
model,
|
||||||
|
|
@ -264,9 +302,8 @@ impl ModelClient for OpenRouterClient {
|
||||||
temperature: None,
|
temperature: None,
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut auth_val =
|
let mut auth_val = reqwest::header::HeaderValue::from_str(&format!("Bearer {}", api_key))
|
||||||
reqwest::header::HeaderValue::from_str(&format!("Bearer {}", self.api_key))
|
.unwrap_or_else(|_| reqwest::header::HeaderValue::from_static(""));
|
||||||
.unwrap_or_else(|_| reqwest::header::HeaderValue::from_static(""));
|
|
||||||
auth_val.set_sensitive(true);
|
auth_val.set_sensitive(true);
|
||||||
|
|
||||||
let res = self
|
let res = self
|
||||||
|
|
@ -283,9 +320,9 @@ impl ModelClient for OpenRouterClient {
|
||||||
let status = res.status();
|
let status = res.status();
|
||||||
let headers = res.headers().clone();
|
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();
|
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();
|
let mut stream = res.bytes_stream();
|
||||||
|
|
@ -353,9 +390,11 @@ mod tests {
|
||||||
reqwest::header::HeaderValue::from_static("12"),
|
reqwest::header::HeaderValue::from_static("12"),
|
||||||
);
|
);
|
||||||
|
|
||||||
|
assert!(OpenRouterClient::check_status(reqwest::StatusCode::OK, "", &headers).is_none());
|
||||||
|
|
||||||
let err_401 =
|
let err_401 =
|
||||||
OpenRouterClient::check_status(reqwest::StatusCode::UNAUTHORIZED, "bad key", &headers)
|
OpenRouterClient::check_status(reqwest::StatusCode::UNAUTHORIZED, "bad key", &headers)
|
||||||
.unwrap_err();
|
.expect("401 must map to an error");
|
||||||
match err_401 {
|
match err_401 {
|
||||||
ClientError::Unauthorized => (),
|
ClientError::Unauthorized => (),
|
||||||
_ => panic!("Expected Unauthorized, got {:?}", err_401),
|
_ => panic!("Expected Unauthorized, got {:?}", err_401),
|
||||||
|
|
@ -366,7 +405,7 @@ mod tests {
|
||||||
"low balance",
|
"low balance",
|
||||||
&headers,
|
&headers,
|
||||||
)
|
)
|
||||||
.unwrap_err();
|
.expect("402 must map to an error");
|
||||||
match err_402 {
|
match err_402 {
|
||||||
ClientError::PaymentRequired => (),
|
ClientError::PaymentRequired => (),
|
||||||
_ => panic!("Expected PaymentRequired, got {:?}", err_402),
|
_ => panic!("Expected PaymentRequired, got {:?}", err_402),
|
||||||
|
|
@ -377,7 +416,7 @@ mod tests {
|
||||||
"rate limited",
|
"rate limited",
|
||||||
&headers,
|
&headers,
|
||||||
)
|
)
|
||||||
.unwrap_err();
|
.expect("429 must map to an error");
|
||||||
match err_429 {
|
match err_429 {
|
||||||
ClientError::RateLimited(secs) => assert_eq!(secs, 12),
|
ClientError::RateLimited(secs) => assert_eq!(secs, 12),
|
||||||
_ => panic!("Expected RateLimited(12), got {:?}", err_429),
|
_ => panic!("Expected RateLimited(12), got {:?}", err_429),
|
||||||
|
|
@ -388,7 +427,7 @@ mod tests {
|
||||||
"outage",
|
"outage",
|
||||||
&headers,
|
&headers,
|
||||||
)
|
)
|
||||||
.unwrap_err();
|
.expect("503 must map to an error");
|
||||||
match err_503 {
|
match err_503 {
|
||||||
ClientError::ServerError { status, .. } => assert_eq!(status, 503),
|
ClientError::ServerError { status, .. } => assert_eq!(status, 503),
|
||||||
_ => panic!("Expected ServerError 503, got {:?}", err_503),
|
_ => panic!("Expected ServerError 503, got {:?}", err_503),
|
||||||
|
|
|
||||||
186
src/config.rs
186
src/config.rs
|
|
@ -9,6 +9,25 @@ const KEYRING_SERVICE: &str = "brokecrew";
|
||||||
const KEYRING_USER: &str = "openrouter_api_key";
|
const KEYRING_USER: &str = "openrouter_api_key";
|
||||||
const DEFAULT_PERSONAS_TOML: &str = include_str!("../config/personas.example.toml");
|
const DEFAULT_PERSONAS_TOML: &str = include_str!("../config/personas.example.toml");
|
||||||
|
|
||||||
|
fn default_persona_str() -> String {
|
||||||
|
"general".to_string()
|
||||||
|
}
|
||||||
|
fn default_router_model() -> String {
|
||||||
|
"liquid/lfm-2.5-2.6b:free".to_string()
|
||||||
|
}
|
||||||
|
fn default_complexity_threshold() -> u8 {
|
||||||
|
7
|
||||||
|
}
|
||||||
|
fn default_credit_tier() -> CreditTier {
|
||||||
|
CreditTier::LowBalance
|
||||||
|
}
|
||||||
|
fn default_base_url() -> String {
|
||||||
|
"https://openrouter.ai/api/v1".to_string()
|
||||||
|
}
|
||||||
|
fn default_theme_str() -> String {
|
||||||
|
"opencode".to_string()
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub struct AppConfigFile {
|
pub struct AppConfigFile {
|
||||||
#[serde(default = "default_persona_str")]
|
#[serde(default = "default_persona_str")]
|
||||||
|
|
@ -23,25 +42,11 @@ pub struct AppConfigFile {
|
||||||
pub credit_tier: CreditTier,
|
pub credit_tier: CreditTier,
|
||||||
#[serde(default = "default_base_url")]
|
#[serde(default = "default_base_url")]
|
||||||
pub openrouter_base_url: String,
|
pub openrouter_base_url: String,
|
||||||
|
#[serde(default = "default_theme_str")]
|
||||||
|
pub theme: String,
|
||||||
pub api_key: Option<String>,
|
pub api_key: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
fn default_persona_str() -> String {
|
|
||||||
"general".to_string()
|
|
||||||
}
|
|
||||||
fn default_router_model() -> String {
|
|
||||||
"meta-llama/llama-3.2-3b-instruct:free".to_string()
|
|
||||||
}
|
|
||||||
fn default_complexity_threshold() -> u8 {
|
|
||||||
7
|
|
||||||
}
|
|
||||||
fn default_credit_tier() -> CreditTier {
|
|
||||||
CreditTier::LowBalance
|
|
||||||
}
|
|
||||||
fn default_base_url() -> String {
|
|
||||||
"https://openrouter.ai/api/v1".to_string()
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Default for AppConfigFile {
|
impl Default for AppConfigFile {
|
||||||
fn default() -> Self {
|
fn default() -> Self {
|
||||||
Self {
|
Self {
|
||||||
|
|
@ -51,6 +56,7 @@ impl Default for AppConfigFile {
|
||||||
allow_stealth: false,
|
allow_stealth: false,
|
||||||
credit_tier: default_credit_tier(),
|
credit_tier: default_credit_tier(),
|
||||||
openrouter_base_url: default_base_url(),
|
openrouter_base_url: default_base_url(),
|
||||||
|
theme: default_theme_str(),
|
||||||
api_key: None,
|
api_key: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -71,6 +77,7 @@ pub struct AppConfig {
|
||||||
pub allow_stealth: bool,
|
pub allow_stealth: bool,
|
||||||
pub credit_tier: CreditTier,
|
pub credit_tier: CreditTier,
|
||||||
pub openrouter_base_url: String,
|
pub openrouter_base_url: String,
|
||||||
|
pub theme: String,
|
||||||
pub api_key: Option<String>,
|
pub api_key: Option<String>,
|
||||||
pub persona_registry: PersonaRegistry,
|
pub persona_registry: PersonaRegistry,
|
||||||
pub pipelines: Vec<PipelineConfig>,
|
pub pipelines: Vec<PipelineConfig>,
|
||||||
|
|
@ -131,6 +138,7 @@ impl AppConfig {
|
||||||
allow_stealth: app_file.allow_stealth,
|
allow_stealth: app_file.allow_stealth,
|
||||||
credit_tier: app_file.credit_tier,
|
credit_tier: app_file.credit_tier,
|
||||||
openrouter_base_url: app_file.openrouter_base_url,
|
openrouter_base_url: app_file.openrouter_base_url,
|
||||||
|
theme: app_file.theme,
|
||||||
api_key,
|
api_key,
|
||||||
persona_registry: PersonaRegistry::new(personas_config.personas),
|
persona_registry: PersonaRegistry::new(personas_config.personas),
|
||||||
pipelines: personas_config.pipelines,
|
pipelines: personas_config.pipelines,
|
||||||
|
|
@ -138,6 +146,107 @@ impl AppConfig {
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Persists the selected theme into `config.toml`
|
||||||
|
pub fn save_theme(config_dir: &Path, theme: &str) -> 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::<AppConfigFile>(&data).unwrap_or_default()
|
||||||
|
} else {
|
||||||
|
AppConfigFile::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
app_file.theme = theme.to_string();
|
||||||
|
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(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Persists the provider base URL into `config.toml`
|
||||||
|
pub fn save_openrouter_base_url(config_dir: &Path, base_url: &str) -> 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::<AppConfigFile>(&data).unwrap_or_default()
|
||||||
|
} else {
|
||||||
|
AppConfigFile::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
app_file.openrouter_base_url = base_url.to_string();
|
||||||
|
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(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Persists the selected router model into `config.toml`
|
||||||
|
pub fn save_router_model(config_dir: &Path, router_model: &str) -> 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::<AppConfigFile>(&data).unwrap_or_default()
|
||||||
|
} else {
|
||||||
|
AppConfigFile::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
app_file.router_model = router_model.to_string();
|
||||||
|
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(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 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::<AppConfigFile>(&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();
|
||||||
|
if trimmed.is_empty() {
|
||||||
|
return "not configured".to_string();
|
||||||
|
}
|
||||||
|
if trimmed.len() <= 8 {
|
||||||
|
return "••••••••".to_string();
|
||||||
|
}
|
||||||
|
let prefix = &trimmed[..trimmed.len().min(7)];
|
||||||
|
format!("{prefix}••••••••")
|
||||||
|
}
|
||||||
|
|
||||||
/// Resolves the API key using precedence:
|
/// Resolves the API key using precedence:
|
||||||
/// 1. CLI argument
|
/// 1. CLI argument
|
||||||
/// 2. `OPENROUTER_API_KEY` environment variable
|
/// 2. `OPENROUTER_API_KEY` environment variable
|
||||||
|
|
@ -237,4 +346,49 @@ mod tests {
|
||||||
assert!(parsed.personas.len() >= 3);
|
assert!(parsed.personas.len() >= 3);
|
||||||
assert!(!parsed.pipelines.is_empty());
|
assert!(!parsed.pipelines.is_empty());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_theme_persistence() {
|
||||||
|
let temp_dir = tempfile::tempdir().unwrap();
|
||||||
|
AppConfig::save_theme(temp_dir.path(), "dracula").expect("save theme");
|
||||||
|
let loaded =
|
||||||
|
AppConfig::load_from_paths(Some(&temp_dir.path().join("config.toml")), None, None)
|
||||||
|
.expect("load config");
|
||||||
|
assert_eq!(loaded.theme, "dracula");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_provider_url_and_router_model_persistence() {
|
||||||
|
let temp_dir = tempfile::tempdir().unwrap();
|
||||||
|
AppConfig::save_openrouter_base_url(temp_dir.path(), "http://localhost:11434/v1")
|
||||||
|
.expect("save base url");
|
||||||
|
AppConfig::save_router_model(temp_dir.path(), "qwen/qwen-2.5-7b:free")
|
||||||
|
.expect("save router model");
|
||||||
|
|
||||||
|
let loaded =
|
||||||
|
AppConfig::load_from_paths(Some(&temp_dir.path().join("config.toml")), None, None)
|
||||||
|
.expect("load config");
|
||||||
|
assert_eq!(loaded.openrouter_base_url, "http://localhost:11434/v1");
|
||||||
|
assert_eq!(loaded.router_model, "qwen/qwen-2.5-7b:free");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_mask_api_key() {
|
||||||
|
assert_eq!(AppConfig::mask_api_key(""), "not configured");
|
||||||
|
assert_eq!(AppConfig::mask_api_key("short"), "••••••••");
|
||||||
|
assert_eq!(
|
||||||
|
AppConfig::mask_api_key("sk-or-v1-abcdef123456"),
|
||||||
|
"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);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
11
src/lib.rs
11
src/lib.rs
|
|
@ -7,11 +7,18 @@ pub mod persona;
|
||||||
pub mod pipeline;
|
pub mod pipeline;
|
||||||
pub mod registry;
|
pub mod registry;
|
||||||
pub mod router;
|
pub mod router;
|
||||||
|
pub mod session;
|
||||||
|
pub mod tool;
|
||||||
|
|
||||||
pub use budget::{BudgetExceeded, BudgetStatus, BudgetTracker, CreditTier};
|
pub use budget::{estimate_tokens, BudgetExceeded, BudgetStatus, BudgetTracker, CreditTier};
|
||||||
pub use client::{ClientError, ModelClient, OpenRouterClient};
|
pub use client::{ClientError, ModelClient, OpenRouterClient};
|
||||||
pub use config::{AppConfig, AppConfigFile, PersonasConfigFile};
|
pub use config::{AppConfig, AppConfigFile, PersonasConfigFile};
|
||||||
pub use persona::{DispatchResult, ModelAttempt, ModelCandidate, Persona, PersonaRegistry, Tier};
|
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 registry::{ModelInfo, ModelPricing, ModelRegistry};
|
||||||
pub use router::{KeywordRouter, LlmRouter, RouteDecision, RouteStage, Router};
|
pub use router::{KeywordRouter, LlmRouter, RouteDecision, RouteStage, Router};
|
||||||
|
pub use session::{Session, SessionMessage, SessionSummary};
|
||||||
|
pub use tool::{ToolEngine, ToolOutput};
|
||||||
|
|
|
||||||
228
src/main.rs
228
src/main.rs
|
|
@ -5,7 +5,9 @@ use brokecrew::budget::BudgetTracker;
|
||||||
use brokecrew::client::OpenRouterClient;
|
use brokecrew::client::OpenRouterClient;
|
||||||
use brokecrew::config::AppConfig;
|
use brokecrew::config::AppConfig;
|
||||||
use brokecrew::persona::Tier;
|
use brokecrew::persona::Tier;
|
||||||
use brokecrew::pipeline::{PipelineConfig, PipelineEvent, PipelineExecutor};
|
use brokecrew::pipeline::{
|
||||||
|
CliRunParams, DispatchParams, PipelineEvent, PipelineExecutor, PipelineRunParams,
|
||||||
|
};
|
||||||
use brokecrew::registry::ModelRegistry;
|
use brokecrew::registry::ModelRegistry;
|
||||||
use brokecrew::router::{RouteStage, Router};
|
use brokecrew::router::{RouteStage, Router};
|
||||||
use clap::Parser;
|
use clap::Parser;
|
||||||
|
|
@ -18,13 +20,39 @@ use std::time::Duration;
|
||||||
use tokio::sync::mpsc;
|
use tokio::sync::mpsc;
|
||||||
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
|
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
|
||||||
|
|
||||||
fn init_logging(json_logs: bool, verbose: bool) {
|
fn init_logging(
|
||||||
|
json_logs: bool,
|
||||||
|
verbose: bool,
|
||||||
|
is_tui: bool,
|
||||||
|
log_path: Option<std::path::PathBuf>,
|
||||||
|
) {
|
||||||
let filter = if verbose {
|
let filter = if verbose {
|
||||||
EnvFilter::new("brokecrew=debug,reqwest=debug")
|
EnvFilter::new("brokecrew=debug,reqwest=debug")
|
||||||
} else {
|
} else {
|
||||||
EnvFilter::new("brokecrew=info")
|
EnvFilter::new("brokecrew=info")
|
||||||
};
|
};
|
||||||
|
|
||||||
|
if is_tui {
|
||||||
|
if let Some(path) = log_path {
|
||||||
|
if let Ok(file) = std::fs::OpenOptions::new()
|
||||||
|
.create(true)
|
||||||
|
.append(true)
|
||||||
|
.open(path)
|
||||||
|
{
|
||||||
|
tracing_subscriber::registry()
|
||||||
|
.with(filter)
|
||||||
|
.with(tracing_subscriber::fmt::layer().with_writer(file))
|
||||||
|
.init();
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
tracing_subscriber::registry()
|
||||||
|
.with(filter)
|
||||||
|
.with(tracing_subscriber::fmt::layer().with_writer(std::io::sink))
|
||||||
|
.init();
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
if json_logs {
|
if json_logs {
|
||||||
tracing_subscriber::registry()
|
tracing_subscriber::registry()
|
||||||
.with(filter)
|
.with(filter)
|
||||||
|
|
@ -69,7 +97,6 @@ fn create_spinner(message: &'static str) -> ProgressBar {
|
||||||
#[tokio::main]
|
#[tokio::main]
|
||||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||||
let cli = Cli::parse();
|
let cli = Cli::parse();
|
||||||
init_logging(cli.json_logs, cli.verbose);
|
|
||||||
|
|
||||||
let config = AppConfig::load_from_paths(
|
let config = AppConfig::load_from_paths(
|
||||||
cli.config_file.as_deref(),
|
cli.config_file.as_deref(),
|
||||||
|
|
@ -78,6 +105,19 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||||
)
|
)
|
||||||
.map_err(|e| format!("Configuration error: {e}"))?;
|
.map_err(|e| format!("Configuration error: {e}"))?;
|
||||||
|
|
||||||
|
let is_tui = match cli.command {
|
||||||
|
Some(Commands::Tui(_)) => true,
|
||||||
|
None => !cli.cli && io::stdout().is_terminal(),
|
||||||
|
_ => false,
|
||||||
|
};
|
||||||
|
|
||||||
|
init_logging(
|
||||||
|
cli.json_logs,
|
||||||
|
cli.verbose,
|
||||||
|
is_tui,
|
||||||
|
Some(config.config_dir.join("brokecrew.log")),
|
||||||
|
);
|
||||||
|
|
||||||
let budget_path = config.config_dir.join("budget.json");
|
let budget_path = config.config_dir.join("budget.json");
|
||||||
let mut budget_tracker = BudgetTracker::new(Some(budget_path), config.credit_tier);
|
let mut budget_tracker = BudgetTracker::new(Some(budget_path), config.credit_tier);
|
||||||
let model_registry = Arc::new(ModelRegistry::new());
|
let model_registry = Arc::new(ModelRegistry::new());
|
||||||
|
|
@ -106,12 +146,69 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||||
.map_err(|e| format!("Failed to save API key to OS credential store: {e}"))?;
|
.map_err(|e| format!("Failed to save API key to OS credential store: {e}"))?;
|
||||||
println!(
|
println!(
|
||||||
"{}",
|
"{}",
|
||||||
"✔ OpenRouter API key saved securely in OS credential store."
|
"✓ OpenRouter API key saved securely in OS credential store."
|
||||||
.bright_green()
|
.bright_green()
|
||||||
.bold()
|
.bold()
|
||||||
);
|
);
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
ConfigSubcommand::SetUrl { url } => {
|
||||||
|
AppConfig::save_openrouter_base_url(&config.config_dir, &url)
|
||||||
|
.map_err(|e| format!("Failed to save provider base URL: {e}"))?;
|
||||||
|
println!(
|
||||||
|
"{}",
|
||||||
|
format!(
|
||||||
|
"✓ Provider base URL updated to '{}' and saved to config.toml",
|
||||||
|
url
|
||||||
|
)
|
||||||
|
.bright_green()
|
||||||
|
.bold()
|
||||||
|
);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
ConfigSubcommand::SetRouterModel { model } => {
|
||||||
|
AppConfig::save_router_model(&config.config_dir, &model)
|
||||||
|
.map_err(|e| format!("Failed to save router model: {e}"))?;
|
||||||
|
println!(
|
||||||
|
"{}",
|
||||||
|
format!(
|
||||||
|
"✓ Stage B router model updated to '{}' and saved to config.toml",
|
||||||
|
model
|
||||||
|
)
|
||||||
|
.bright_green()
|
||||||
|
.bold()
|
||||||
|
);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
ConfigSubcommand::Show => {
|
||||||
|
let key_str = config.api_key.as_deref().unwrap_or("");
|
||||||
|
let masked = AppConfig::mask_api_key(key_str);
|
||||||
|
println!("\n{}", "BrokeCrew Configuration:".bold().underline());
|
||||||
|
println!(
|
||||||
|
" Provider Base URL: {}",
|
||||||
|
config.openrouter_base_url.bright_cyan()
|
||||||
|
);
|
||||||
|
println!(
|
||||||
|
" Router Model: {}",
|
||||||
|
config.router_model.bright_white().bold()
|
||||||
|
);
|
||||||
|
println!(
|
||||||
|
" Default Persona: {}",
|
||||||
|
config.default_persona.bright_yellow()
|
||||||
|
);
|
||||||
|
println!(" API Key: {}", masked);
|
||||||
|
println!(" Stealth Allowed: {}", config.allow_stealth);
|
||||||
|
println!(" Complexity Threshold: {}", config.complexity_threshold);
|
||||||
|
println!(" Theme: {}", config.theme.bright_magenta());
|
||||||
|
println!(
|
||||||
|
" Config Directory: {}\n",
|
||||||
|
config.config_dir.display().to_string().dimmed()
|
||||||
|
);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
},
|
},
|
||||||
|
|
||||||
Some(Commands::Personas(personas_args)) => match personas_args.action {
|
Some(Commands::Personas(personas_args)) => match personas_args.action {
|
||||||
|
|
@ -225,7 +322,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||||
);
|
);
|
||||||
println!(" Daily requests today: {}", daily_color);
|
println!(" Daily requests today: {}", daily_color);
|
||||||
if status.approaching_daily_limit {
|
if status.approaching_daily_limit {
|
||||||
println!(" {}", "⚠ Notice: Approaching daily free-tier limit. Non-urgent calls may be deferred.".yellow().bold());
|
println!(" {}", "! Notice: Approaching daily free-tier limit. Non-urgent calls may be deferred.".yellow().bold());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
|
|
@ -285,7 +382,7 @@ fn ensure_api_key(config: &AppConfig) -> Result<String, Box<dyn std::error::Erro
|
||||||
if io::stdin().is_terminal() {
|
if io::stdin().is_terminal() {
|
||||||
eprintln!(
|
eprintln!(
|
||||||
"{}",
|
"{}",
|
||||||
"⚡ Welcome to BrokeCrew! No OpenRouter API key found in keyring, env, or config."
|
":: Welcome to BrokeCrew! No OpenRouter API key found in keyring, env, or config."
|
||||||
.bright_cyan()
|
.bright_cyan()
|
||||||
.bold()
|
.bold()
|
||||||
);
|
);
|
||||||
|
|
@ -301,7 +398,7 @@ fn ensure_api_key(config: &AppConfig) -> Result<String, Box<dyn std::error::Erro
|
||||||
let _ = AppConfig::save_api_key_to_keyring(&trimmed);
|
let _ = AppConfig::save_api_key_to_keyring(&trimmed);
|
||||||
eprintln!(
|
eprintln!(
|
||||||
"{}",
|
"{}",
|
||||||
"✔ OpenRouter API key saved securely in OS credential store."
|
"✓ OpenRouter API key saved securely in OS credential store."
|
||||||
.bright_green()
|
.bright_green()
|
||||||
.bold()
|
.bold()
|
||||||
);
|
);
|
||||||
|
|
@ -327,12 +424,11 @@ async fn run_chat(
|
||||||
|
|
||||||
let client = OpenRouterClient::with_base_url(api_key, config.openrouter_base_url.clone());
|
let client = OpenRouterClient::with_base_url(api_key, config.openrouter_base_url.clone());
|
||||||
|
|
||||||
// Start background sync for model registry
|
// Start background sync for model registry (immediate + every 30 min).
|
||||||
let reg_clone = model_registry.clone();
|
let _registry_sync = model_registry.start_background_sync(
|
||||||
let base_url = config.openrouter_base_url.clone();
|
config.openrouter_base_url.clone(),
|
||||||
tokio::spawn(async move {
|
std::time::Duration::from_secs(30 * 60),
|
||||||
let _ = reg_clone.fetch_live(&base_url).await;
|
);
|
||||||
});
|
|
||||||
|
|
||||||
let router = Router::new(
|
let router = Router::new(
|
||||||
config.default_persona.clone(),
|
config.default_persona.clone(),
|
||||||
|
|
@ -410,13 +506,13 @@ async fn execute_prompt(
|
||||||
router: &Router,
|
router: &Router,
|
||||||
client: &OpenRouterClient,
|
client: &OpenRouterClient,
|
||||||
) -> Result<(), Box<dyn std::error::Error>> {
|
) -> Result<(), Box<dyn std::error::Error>> {
|
||||||
// 1. Budget check
|
// 1. Budget pre-check (advisory only; each dispatch records atomically).
|
||||||
match budget_tracker.check_budget() {
|
match budget_tracker.check_budget() {
|
||||||
Ok(status) => {
|
Ok(status) => {
|
||||||
if status.approaching_daily_limit {
|
if status.approaching_daily_limit {
|
||||||
eprintln!(
|
eprintln!(
|
||||||
"{} Daily usage: {}/{}",
|
"{} Daily usage: {}/{}",
|
||||||
"⚠ Budget notice: Approaching daily ceiling.".yellow(),
|
"! Budget notice: Approaching daily ceiling.".yellow(),
|
||||||
status.daily_requests,
|
status.daily_requests,
|
||||||
status.daily_limit
|
status.daily_limit
|
||||||
);
|
);
|
||||||
|
|
@ -439,7 +535,7 @@ async fn execute_prompt(
|
||||||
.find(|p| p.name == *pipe_name)
|
.find(|p| p.name == *pipe_name)
|
||||||
.ok_or_else(|| format!("Pipeline '{pipe_name}' not found in configuration"))?;
|
.ok_or_else(|| format!("Pipeline '{pipe_name}' not found in configuration"))?;
|
||||||
|
|
||||||
run_pipeline(
|
run_pipeline(CliRunParams {
|
||||||
pipeline,
|
pipeline,
|
||||||
config,
|
config,
|
||||||
budget_tracker,
|
budget_tracker,
|
||||||
|
|
@ -448,7 +544,7 @@ async fn execute_prompt(
|
||||||
prompt,
|
prompt,
|
||||||
allow_stealth,
|
allow_stealth,
|
||||||
require_paid,
|
require_paid,
|
||||||
)
|
})
|
||||||
.await?;
|
.await?;
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
|
|
@ -467,7 +563,7 @@ async fn execute_prompt(
|
||||||
match decision {
|
match decision {
|
||||||
Ok(dec) => {
|
Ok(dec) => {
|
||||||
if dec.stage_used == RouteStage::StageBLlm {
|
if dec.stage_used == RouteStage::StageBLlm {
|
||||||
let _ = budget_tracker.record_call(false);
|
let _ = budget_tracker.try_acquire(false);
|
||||||
}
|
}
|
||||||
let info = match dec.stage_used {
|
let info = match dec.stage_used {
|
||||||
RouteStage::StageAKeyword => "routed via keyword pre-filter".to_string(),
|
RouteStage::StageAKeyword => "routed via keyword pre-filter".to_string(),
|
||||||
|
|
@ -488,7 +584,7 @@ async fn execute_prompt(
|
||||||
if let Some(info) = route_info {
|
if let Some(info) = route_info {
|
||||||
eprintln!(
|
eprintln!(
|
||||||
"{}",
|
"{}",
|
||||||
format!("⚡ [BrokeCrew: {info} ➔ {}]", target_persona).dimmed()
|
format!("› [BrokeCrew: {info} ➔ {}]", target_persona).dimmed()
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -496,16 +592,16 @@ async fn execute_prompt(
|
||||||
let persona_tag = get_persona_color_tag(&target_persona);
|
let persona_tag = get_persona_color_tag(&target_persona);
|
||||||
let spinner = create_spinner("Thinking...");
|
let spinner = create_spinner("Thinking...");
|
||||||
|
|
||||||
let dispatch_future = PipelineExecutor::dispatch_persona(
|
let dispatch_future = PipelineExecutor::dispatch_persona(DispatchParams {
|
||||||
&target_persona,
|
persona_name: &target_persona,
|
||||||
&config.persona_registry,
|
persona_registry: &config.persona_registry,
|
||||||
model_registry,
|
model_registry,
|
||||||
client,
|
client,
|
||||||
prompt,
|
prompt,
|
||||||
allow_stealth,
|
allow_stealth,
|
||||||
require_paid,
|
require_paid,
|
||||||
Some(chunk_tx),
|
chunk_sender: Some(chunk_tx),
|
||||||
);
|
});
|
||||||
|
|
||||||
let printer = tokio::spawn(async move {
|
let printer = tokio::spawn(async move {
|
||||||
let mut first = true;
|
let mut first = true;
|
||||||
|
|
@ -529,14 +625,14 @@ async fn execute_prompt(
|
||||||
if res.fell_back {
|
if res.fell_back {
|
||||||
eprintln!(
|
eprintln!(
|
||||||
"{}",
|
"{}",
|
||||||
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
|
// Record budget for each failed attempt plus the successful one
|
||||||
for _ in 0..res.attempts.len() {
|
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) => {
|
Err(e) => {
|
||||||
eprintln!("\n{} {}", "Error dispatching to persona:".red().bold(), e);
|
eprintln!("\n{} {}", "Error dispatching to persona:".red().bold(), e);
|
||||||
|
|
@ -546,36 +642,36 @@ async fn execute_prompt(
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
#[allow(clippy::too_many_arguments)]
|
async fn run_pipeline(params: CliRunParams<'_>) -> Result<(), Box<dyn std::error::Error>> {
|
||||||
async fn run_pipeline(
|
let CliRunParams {
|
||||||
pipeline: &PipelineConfig,
|
pipeline,
|
||||||
config: &AppConfig,
|
config,
|
||||||
budget_tracker: &mut BudgetTracker,
|
budget_tracker,
|
||||||
model_registry: &ModelRegistry,
|
model_registry,
|
||||||
client: &OpenRouterClient,
|
client,
|
||||||
prompt: &str,
|
prompt,
|
||||||
allow_stealth: bool,
|
allow_stealth,
|
||||||
require_paid: bool,
|
require_paid,
|
||||||
) -> Result<(), Box<dyn std::error::Error>> {
|
} = params;
|
||||||
println!(
|
println!(
|
||||||
"\n{} Starting pipeline '{}' ({} steps)",
|
"\n{} Starting pipeline '{}' ({} steps)",
|
||||||
"⚡".bright_cyan(),
|
"›".bright_cyan(),
|
||||||
pipeline.name.bold(),
|
pipeline.name.bold(),
|
||||||
pipeline.steps.len()
|
pipeline.steps.len()
|
||||||
);
|
);
|
||||||
|
|
||||||
let (event_tx, mut event_rx) = mpsc::channel::<PipelineEvent>(100);
|
let (event_tx, mut event_rx) = mpsc::channel::<PipelineEvent>(100);
|
||||||
|
|
||||||
let exec_future = PipelineExecutor::execute(
|
let exec_future = PipelineExecutor::execute(PipelineRunParams {
|
||||||
pipeline,
|
pipeline,
|
||||||
&config.persona_registry,
|
persona_registry: &config.persona_registry,
|
||||||
model_registry,
|
model_registry,
|
||||||
client,
|
client,
|
||||||
prompt,
|
initial_prompt: prompt,
|
||||||
allow_stealth,
|
allow_stealth,
|
||||||
require_paid,
|
require_paid,
|
||||||
Some(event_tx),
|
event_sender: Some(event_tx),
|
||||||
);
|
});
|
||||||
|
|
||||||
let event_handler = tokio::spawn(async move {
|
let event_handler = tokio::spawn(async move {
|
||||||
let mut current_spinner: Option<ProgressBar> = None;
|
let mut current_spinner: Option<ProgressBar> = None;
|
||||||
|
|
@ -586,23 +682,46 @@ async fn run_pipeline(
|
||||||
step_index,
|
step_index,
|
||||||
total_steps,
|
total_steps,
|
||||||
persona,
|
persona,
|
||||||
model,
|
|
||||||
..
|
|
||||||
} => {
|
} => {
|
||||||
if let Some(sp) = current_spinner.take() {
|
if let Some(sp) = current_spinner.take() {
|
||||||
sp.finish_and_clear();
|
sp.finish_and_clear();
|
||||||
}
|
}
|
||||||
let badge = get_persona_color_tag(&persona);
|
let badge = get_persona_color_tag(&persona);
|
||||||
println!(
|
println!(
|
||||||
"\n--- Step {}/{} {} ({}) ---",
|
"\n--- Step {}/{} {} ---",
|
||||||
step_index + 1,
|
step_index + 1,
|
||||||
total_steps,
|
total_steps,
|
||||||
badge,
|
badge,
|
||||||
model.dimmed()
|
|
||||||
);
|
);
|
||||||
let sp = create_spinner("Thinking...");
|
let sp = create_spinner("Thinking...");
|
||||||
current_spinner = Some(sp);
|
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,
|
||||||
|
total_candidates,
|
||||||
|
..
|
||||||
|
} => {
|
||||||
|
if let Some(ref sp) = current_spinner {
|
||||||
|
sp.set_message(format!("Trying {model} ({attempt}/{total_candidates})..."));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
PipelineEvent::CandidateFailed { model, error, .. } => {
|
||||||
|
tracing::debug!(model = %model, error = %error, "Candidate attempt failed");
|
||||||
|
}
|
||||||
PipelineEvent::Chunk { delta, .. } => {
|
PipelineEvent::Chunk { delta, .. } => {
|
||||||
if let Some(sp) = current_spinner.take() {
|
if let Some(sp) = current_spinner.take() {
|
||||||
sp.finish_and_clear();
|
sp.finish_and_clear();
|
||||||
|
|
@ -610,13 +729,18 @@ async fn run_pipeline(
|
||||||
print!("{delta}");
|
print!("{delta}");
|
||||||
let _ = io::stdout().flush();
|
let _ = io::stdout().flush();
|
||||||
}
|
}
|
||||||
PipelineEvent::StepCompleted { fell_back, .. } => {
|
PipelineEvent::StepCompleted {
|
||||||
|
model_used,
|
||||||
|
fell_back,
|
||||||
|
..
|
||||||
|
} => {
|
||||||
if let Some(sp) = current_spinner.take() {
|
if let Some(sp) = current_spinner.take() {
|
||||||
sp.finish_and_clear();
|
sp.finish_and_clear();
|
||||||
}
|
}
|
||||||
println!();
|
println!();
|
||||||
|
eprintln!("{}", format!("› served by {model_used}").dimmed());
|
||||||
if fell_back {
|
if fell_back {
|
||||||
eprintln!("{}", "⚠ Fallback occurred during step execution".yellow());
|
eprintln!("{}", "! Fallback occurred during step execution".yellow());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -628,9 +752,9 @@ async fn run_pipeline(
|
||||||
|
|
||||||
for r in &results {
|
for r in &results {
|
||||||
for _ in 0..r.attempts.len() {
|
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(())
|
Ok(())
|
||||||
|
|
|
||||||
|
|
@ -102,6 +102,21 @@ impl PersonaRegistry {
|
||||||
Ok(eligible)
|
Ok(eligible)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Returns a new PersonaRegistry with tool execution prompts injected into persona system prompts
|
||||||
|
pub fn with_tool_capabilities(&self) -> Self {
|
||||||
|
let tool_prompt = crate::tool::ToolEngine::system_tool_prompt();
|
||||||
|
let personas = self
|
||||||
|
.personas
|
||||||
|
.iter()
|
||||||
|
.map(|p| {
|
||||||
|
let mut cp = p.clone();
|
||||||
|
cp.system_prompt.push_str(tool_prompt);
|
||||||
|
cp
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
Self { personas }
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
|
|
|
||||||
457
src/pipeline.rs
457
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::persona::{DispatchResult, ModelAttempt, ModelCandidate, PersonaRegistry, Tier};
|
||||||
use crate::registry::ModelRegistry;
|
use crate::registry::ModelRegistry;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
@ -17,6 +19,23 @@ pub struct PipelineStep {
|
||||||
pub struct PipelineConfig {
|
pub struct PipelineConfig {
|
||||||
pub name: String,
|
pub name: String,
|
||||||
pub steps: Vec<PipelineStep>,
|
pub steps: Vec<PipelineStep>,
|
||||||
|
/// 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<DispatchResult>,
|
||||||
|
pub failed_step: Option<usize>,
|
||||||
|
pub error: Option<ClientError>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
|
|
@ -25,9 +44,21 @@ pub enum PipelineEvent {
|
||||||
step_index: usize,
|
step_index: usize,
|
||||||
total_steps: usize,
|
total_steps: usize,
|
||||||
persona: String,
|
persona: String,
|
||||||
|
},
|
||||||
|
CandidateAttempt {
|
||||||
|
step_index: usize,
|
||||||
|
persona: String,
|
||||||
model: String,
|
model: String,
|
||||||
|
attempt: usize,
|
||||||
|
total_candidates: usize,
|
||||||
tier: Tier,
|
tier: Tier,
|
||||||
},
|
},
|
||||||
|
CandidateFailed {
|
||||||
|
step_index: usize,
|
||||||
|
persona: String,
|
||||||
|
model: String,
|
||||||
|
error: String,
|
||||||
|
},
|
||||||
Chunk {
|
Chunk {
|
||||||
step_index: usize,
|
step_index: usize,
|
||||||
delta: String,
|
delta: String,
|
||||||
|
|
@ -35,27 +66,83 @@ pub enum PipelineEvent {
|
||||||
StepCompleted {
|
StepCompleted {
|
||||||
step_index: usize,
|
step_index: usize,
|
||||||
persona: String,
|
persona: String,
|
||||||
|
model_used: String,
|
||||||
|
tier_used: Tier,
|
||||||
full_reply: String,
|
full_reply: String,
|
||||||
fell_back: bool,
|
fell_back: bool,
|
||||||
attempts: Vec<ModelAttempt>,
|
attempts: Vec<ModelAttempt>,
|
||||||
},
|
},
|
||||||
|
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<mpsc::Sender<String>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 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<mpsc::Sender<PipelineEvent>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 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;
|
pub struct PipelineExecutor;
|
||||||
|
|
||||||
impl PipelineExecutor {
|
impl PipelineExecutor {
|
||||||
/// Executes a single persona query with candidate fallback walk
|
/// Executes a single persona query with candidate fallback walk
|
||||||
#[allow(clippy::too_many_arguments)]
|
|
||||||
pub async fn dispatch_persona(
|
pub async fn dispatch_persona(
|
||||||
persona_name: &str,
|
params: DispatchParams<'_>,
|
||||||
persona_registry: &PersonaRegistry,
|
|
||||||
model_registry: &ModelRegistry,
|
|
||||||
client: &dyn ModelClient,
|
|
||||||
prompt: &str,
|
|
||||||
allow_stealth: bool,
|
|
||||||
require_paid: bool,
|
|
||||||
chunk_sender: Option<mpsc::Sender<String>>,
|
|
||||||
) -> Result<DispatchResult, ClientError> {
|
) -> Result<DispatchResult, ClientError> {
|
||||||
|
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<mpsc::Sender<PipelineEvent>>,
|
||||||
|
) -> Result<DispatchResult, ClientError> {
|
||||||
|
let DispatchParams {
|
||||||
|
persona_name,
|
||||||
|
persona_registry,
|
||||||
|
model_registry,
|
||||||
|
client,
|
||||||
|
prompt,
|
||||||
|
allow_stealth,
|
||||||
|
require_paid,
|
||||||
|
chunk_sender,
|
||||||
|
} = params;
|
||||||
let persona = persona_registry.get(persona_name).ok_or_else(|| {
|
let persona = persona_registry.get(persona_name).ok_or_else(|| {
|
||||||
ClientError::ModelUnavailable(format!("Persona '{persona_name}' not found"))
|
ClientError::ModelUnavailable(format!("Persona '{persona_name}' not found"))
|
||||||
})?;
|
})?;
|
||||||
|
|
@ -112,6 +199,19 @@ impl PipelineExecutor {
|
||||||
"Attempting candidate dispatch"
|
"Attempting candidate dispatch"
|
||||||
);
|
);
|
||||||
|
|
||||||
|
if let Some(ref tx) = event_sender {
|
||||||
|
let _ = tx
|
||||||
|
.send(PipelineEvent::CandidateAttempt {
|
||||||
|
step_index: 0,
|
||||||
|
persona: persona_name.to_string(),
|
||||||
|
model: model_id.clone(),
|
||||||
|
attempt: idx + 1,
|
||||||
|
total_candidates,
|
||||||
|
tier: candidate.tier,
|
||||||
|
})
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
// Channel for streaming chunks if sender provided
|
// Channel for streaming chunks if sender provided
|
||||||
let (tx, mut rx) = mpsc::channel::<String>(100);
|
let (tx, mut rx) = mpsc::channel::<String>(100);
|
||||||
let forwarder = {
|
let forwarder = {
|
||||||
|
|
@ -165,6 +265,17 @@ impl PipelineExecutor {
|
||||||
error: Some(err_msg.clone()),
|
error: Some(err_msg.clone()),
|
||||||
});
|
});
|
||||||
|
|
||||||
|
if let Some(ref tx) = event_sender {
|
||||||
|
let _ = tx
|
||||||
|
.send(PipelineEvent::CandidateFailed {
|
||||||
|
step_index: 0,
|
||||||
|
persona: persona_name.to_string(),
|
||||||
|
model: model_id.clone(),
|
||||||
|
error: err_msg.clone(),
|
||||||
|
})
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
// Fast-fail on auth or billing errors
|
// Fast-fail on auth or billing errors
|
||||||
match e {
|
match e {
|
||||||
ClientError::Unauthorized | ClientError::PaymentRequired => return Err(e),
|
ClientError::Unauthorized | ClientError::PaymentRequired => return Err(e),
|
||||||
|
|
@ -209,26 +320,55 @@ impl PipelineExecutor {
|
||||||
)))
|
)))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Executes an entire pipeline sequentially, passing outputs between steps
|
/// Executes an entire pipeline sequentially, passing outputs between steps.
|
||||||
#[allow(clippy::too_many_arguments)]
|
/// 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(
|
pub async fn execute(
|
||||||
pipeline: &PipelineConfig,
|
params: PipelineRunParams<'_>,
|
||||||
persona_registry: &PersonaRegistry,
|
|
||||||
model_registry: &ModelRegistry,
|
|
||||||
client: &dyn ModelClient,
|
|
||||||
initial_prompt: &str,
|
|
||||||
allow_stealth: bool,
|
|
||||||
require_paid: bool,
|
|
||||||
event_sender: Option<mpsc::Sender<PipelineEvent>>,
|
|
||||||
) -> Result<Vec<DispatchResult>, ClientError> {
|
) -> Result<Vec<DispatchResult>, 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 results = Vec::new();
|
||||||
let mut previous_output = String::new();
|
let mut previous_output = String::new();
|
||||||
let total_steps = pipeline.steps.len();
|
let total_steps = pipeline.steps.len();
|
||||||
|
|
||||||
for (step_idx, step) in pipeline.steps.iter().enumerate() {
|
for (step_idx, step) in pipeline.steps.iter().enumerate() {
|
||||||
let persona = persona_registry.get(&step.persona).ok_or_else(|| {
|
if persona_registry.get(&step.persona).is_none() {
|
||||||
ClientError::ModelUnavailable(format!("Persona '{}' not found", step.persona))
|
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
|
// Prepare prompt for this step
|
||||||
let step_prompt = if step.pass_previous_output && !previous_output.is_empty() {
|
let step_prompt = if step.pass_previous_output && !previous_output.is_empty() {
|
||||||
|
|
@ -262,58 +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 {
|
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
|
let _ = tx
|
||||||
.send(PipelineEvent::StepStarted {
|
.send(PipelineEvent::StepStarted {
|
||||||
step_index: step_idx,
|
step_index: step_idx,
|
||||||
total_steps,
|
total_steps,
|
||||||
persona: step.persona.clone(),
|
persona: step.persona.clone(),
|
||||||
model: top_candidate.0,
|
|
||||||
tier: top_candidate.1,
|
|
||||||
})
|
})
|
||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
|
|
||||||
let dispatch_result = Self::dispatch_persona(
|
let dispatch_result = Self::dispatch_persona_with_events(
|
||||||
&step.persona,
|
DispatchParams {
|
||||||
persona_registry,
|
persona_name: &step.persona,
|
||||||
model_registry,
|
persona_registry,
|
||||||
client,
|
model_registry,
|
||||||
&step_prompt,
|
client,
|
||||||
allow_stealth,
|
prompt: &step_prompt,
|
||||||
require_paid,
|
allow_stealth,
|
||||||
Some(chunk_tx),
|
require_paid,
|
||||||
|
chunk_sender: Some(chunk_tx),
|
||||||
|
},
|
||||||
|
event_sender.clone(),
|
||||||
)
|
)
|
||||||
.await?;
|
.await;
|
||||||
|
|
||||||
let _ = forwarder.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 {
|
if let Some(ref tx) = event_sender {
|
||||||
let _ = tx
|
let _ = tx
|
||||||
.send(PipelineEvent::StepCompleted {
|
.send(PipelineEvent::StepCompleted {
|
||||||
step_index: step_idx,
|
step_index: step_idx,
|
||||||
persona: step.persona.clone(),
|
persona: step.persona.clone(),
|
||||||
full_reply: dispatch_result.reply.clone(),
|
model_used: dispatch_result.model_used.clone(),
|
||||||
fell_back: dispatch_result.fell_back,
|
tier_used: dispatch_result.tier_used,
|
||||||
attempts: dispatch_result.attempts.clone(),
|
full_reply: dispatch_result.reply.clone(),
|
||||||
})
|
fell_back: dispatch_result.fell_back,
|
||||||
.await;
|
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,
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -395,6 +561,7 @@ mod tests {
|
||||||
custom_prompt_template: None,
|
custom_prompt_template: None,
|
||||||
},
|
},
|
||||||
],
|
],
|
||||||
|
abort_on_failure: true,
|
||||||
};
|
};
|
||||||
|
|
||||||
let calls = Arc::new(AtomicUsize::new(0));
|
let calls = Arc::new(AtomicUsize::new(0));
|
||||||
|
|
@ -402,16 +569,16 @@ mod tests {
|
||||||
calls: calls.clone(),
|
calls: calls.clone(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let results = PipelineExecutor::execute(
|
let results = PipelineExecutor::execute(PipelineRunParams {
|
||||||
&pipeline,
|
pipeline: &pipeline,
|
||||||
&persona_reg,
|
persona_registry: &persona_reg,
|
||||||
&model_reg,
|
model_registry: &model_reg,
|
||||||
&client,
|
client: &client,
|
||||||
"Build website",
|
initial_prompt: "Build website",
|
||||||
false,
|
allow_stealth: false,
|
||||||
false,
|
require_paid: false,
|
||||||
None,
|
event_sender: None,
|
||||||
)
|
})
|
||||||
.await
|
.await
|
||||||
.expect("pipeline execution failed");
|
.expect("pipeline execution failed");
|
||||||
|
|
||||||
|
|
@ -424,4 +591,160 @@ mod tests {
|
||||||
assert!(results[1].reply.contains("Output from coder-model"));
|
assert!(results[1].reply.contains("Output from coder-model"));
|
||||||
assert_eq!(calls.load(Ordering::SeqCst), 2);
|
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<String, ClientError> {
|
||||||
|
Err(ClientError::ServerError {
|
||||||
|
status: 503,
|
||||||
|
message: "down".into(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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<String, ClientError> {
|
||||||
|
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<String>,
|
||||||
|
) -> Result<String, ClientError> {
|
||||||
|
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::<PipelineEvent>(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());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -33,6 +33,7 @@ struct OpenRouterModelsResponse {
|
||||||
pub struct ModelRegistry {
|
pub struct ModelRegistry {
|
||||||
models: Arc<RwLock<HashMap<String, ModelInfo>>>,
|
models: Arc<RwLock<HashMap<String, ModelInfo>>>,
|
||||||
last_refreshed: Arc<RwLock<Option<Instant>>>,
|
last_refreshed: Arc<RwLock<Option<Instant>>>,
|
||||||
|
http: reqwest::Client,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for ModelRegistry {
|
impl Default for ModelRegistry {
|
||||||
|
|
@ -43,9 +44,14 @@ impl Default for ModelRegistry {
|
||||||
|
|
||||||
impl ModelRegistry {
|
impl ModelRegistry {
|
||||||
pub fn new() -> Self {
|
pub fn new() -> Self {
|
||||||
|
let http = reqwest::Client::builder()
|
||||||
|
.timeout(Duration::from_secs(15))
|
||||||
|
.build()
|
||||||
|
.unwrap_or_else(|_| reqwest::Client::new());
|
||||||
Self {
|
Self {
|
||||||
models: Arc::new(RwLock::new(HashMap::new())),
|
models: Arc::new(RwLock::new(HashMap::new())),
|
||||||
last_refreshed: Arc::new(RwLock::new(None)),
|
last_refreshed: Arc::new(RwLock::new(None)),
|
||||||
|
http,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -84,12 +90,9 @@ impl ModelRegistry {
|
||||||
/// Fetches live models catalog from OpenRouter `/v1/models`
|
/// Fetches live models catalog from OpenRouter `/v1/models`
|
||||||
pub async fn fetch_live(&self, base_url: &str) -> Result<usize, String> {
|
pub async fn fetch_live(&self, base_url: &str) -> Result<usize, String> {
|
||||||
let endpoint = format!("{}/models", base_url.trim_end_matches('/'));
|
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)
|
.get(&endpoint)
|
||||||
.header("HTTP-Referer", "https://github.com/brokecrew/brokecrew")
|
.header("HTTP-Referer", "https://github.com/brokecrew/brokecrew")
|
||||||
.header("X-Title", "BrokeCrew")
|
.header("X-Title", "BrokeCrew")
|
||||||
|
|
@ -112,19 +115,32 @@ impl ModelRegistry {
|
||||||
self.update_from_json(&body).await
|
self.update_from_json(&body).await
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Starts a background task refreshing the models registry periodically
|
/// Starts a background task refreshing the models registry periodically.
|
||||||
pub fn start_background_sync(self: Arc<Self>, base_url: String, interval: Duration) {
|
/// 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<Self>,
|
||||||
|
base_url: String,
|
||||||
|
interval: Duration,
|
||||||
|
) -> tokio::task::JoinHandle<()> {
|
||||||
|
let this = Arc::clone(self);
|
||||||
tokio::spawn(async move {
|
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);
|
let mut ticker = tokio::time::interval(interval);
|
||||||
|
// Skip the immediate first tick: we just synced above.
|
||||||
|
ticker.tick().await;
|
||||||
loop {
|
loop {
|
||||||
ticker.tick().await;
|
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);
|
tracing::warn!("Periodic models registry sync failed: {}", e);
|
||||||
} else {
|
} else {
|
||||||
tracing::debug!("Periodic models registry sync succeeded");
|
tracing::debug!("Periodic models registry sync succeeded");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
});
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn is_model_available(&self, model_id: &str) -> bool {
|
pub async fn is_model_available(&self, model_id: &str) -> bool {
|
||||||
|
|
|
||||||
|
|
@ -23,19 +23,39 @@ pub struct RouteDecision {
|
||||||
|
|
||||||
pub struct Router {
|
pub struct Router {
|
||||||
pub default_persona: String,
|
pub default_persona: String,
|
||||||
pub router_model: String,
|
router_model: arc_swap::ArcSwap<String>,
|
||||||
pub complexity_threshold: u8,
|
complexity_threshold: std::sync::atomic::AtomicU8,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Router {
|
impl Router {
|
||||||
pub fn new(default_persona: String, router_model: String, complexity_threshold: u8) -> Self {
|
pub fn new(default_persona: String, router_model: String, complexity_threshold: u8) -> Self {
|
||||||
Self {
|
Self {
|
||||||
default_persona,
|
default_persona,
|
||||||
router_model,
|
router_model: arc_swap::ArcSwap::new(std::sync::Arc::new(router_model)),
|
||||||
complexity_threshold,
|
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) {
|
||||||
|
self.router_model.store(std::sync::Arc::new(model));
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn get_router_model(&self) -> String {
|
||||||
|
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(
|
pub async fn route(
|
||||||
&self,
|
&self,
|
||||||
registry: &PersonaRegistry,
|
registry: &PersonaRegistry,
|
||||||
|
|
@ -44,7 +64,7 @@ impl Router {
|
||||||
force_escalate: bool,
|
force_escalate: bool,
|
||||||
) -> Result<RouteDecision, ClientError> {
|
) -> Result<RouteDecision, ClientError> {
|
||||||
let complexity_score = LlmRouter::estimate_complexity(prompt);
|
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
|
// Stage A: Free, instant keyword matching
|
||||||
if let Some(matched) = KeywordRouter::classify(registry.list(), prompt) {
|
if let Some(matched) = KeywordRouter::classify(registry.list(), prompt) {
|
||||||
|
|
@ -58,14 +78,10 @@ impl Router {
|
||||||
|
|
||||||
// Stage B: LLM classification fallback
|
// Stage B: LLM classification fallback
|
||||||
let known: Vec<String> = registry.list().iter().map(|p| p.name.clone()).collect();
|
let known: Vec<String> = registry.list().iter().map(|p| p.name.clone()).collect();
|
||||||
let selected = LlmRouter::classify(
|
let router_model = self.get_router_model();
|
||||||
client,
|
let selected =
|
||||||
&self.router_model,
|
LlmRouter::classify(client, &router_model, &known, &self.default_persona, prompt)
|
||||||
&known,
|
.await?;
|
||||||
&self.default_persona,
|
|
||||||
prompt,
|
|
||||||
)
|
|
||||||
.await?;
|
|
||||||
|
|
||||||
Ok(RouteDecision {
|
Ok(RouteDecision {
|
||||||
persona: selected,
|
persona: selected,
|
||||||
|
|
@ -147,4 +163,16 @@ mod tests {
|
||||||
assert_eq!(dec2.persona, "coder");
|
assert_eq!(dec2.persona, "coder");
|
||||||
assert!(dec2.escalate_to_paid);
|
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);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
180
src/session.rs
Normal file
180
src/session.rs
Normal file
|
|
@ -0,0 +1,180 @@
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::fs;
|
||||||
|
use std::io;
|
||||||
|
use std::path::Path;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct SessionMessage {
|
||||||
|
pub sender: String,
|
||||||
|
pub content: String,
|
||||||
|
pub model: Option<String>,
|
||||||
|
pub timestamp: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct Session {
|
||||||
|
pub id: String,
|
||||||
|
pub title: String,
|
||||||
|
pub persona: String,
|
||||||
|
pub created_at: String,
|
||||||
|
pub updated_at: String,
|
||||||
|
pub messages: Vec<SessionMessage>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct SessionSummary {
|
||||||
|
pub id: String,
|
||||||
|
pub title: String,
|
||||||
|
pub persona: String,
|
||||||
|
pub updated_at: String,
|
||||||
|
pub message_count: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Session {
|
||||||
|
pub fn new(persona: &str) -> Self {
|
||||||
|
let now: DateTime<Utc> = Utc::now();
|
||||||
|
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::<u32>() % 0xFFFF;
|
||||||
|
let id = format!("session-{timestamp_str}-{rand_suffix:04x}");
|
||||||
|
|
||||||
|
Self {
|
||||||
|
id,
|
||||||
|
title: "New Session".to_string(),
|
||||||
|
persona: persona.to_string(),
|
||||||
|
created_at: now.to_rfc3339(),
|
||||||
|
updated_at: now.to_rfc3339(),
|
||||||
|
messages: Vec::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn add_message(&mut self, sender: &str, content: &str, model: Option<&str>) {
|
||||||
|
let now: DateTime<Utc> = Utc::now();
|
||||||
|
self.updated_at = now.to_rfc3339();
|
||||||
|
|
||||||
|
// Update session title from first user prompt if still default
|
||||||
|
if self.title == "New Session" && sender.eq_ignore_ascii_case("you") {
|
||||||
|
let truncated: String = content.chars().take(40).collect();
|
||||||
|
self.title = truncated.trim().to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
self.messages.push(SessionMessage {
|
||||||
|
sender: sender.to_string(),
|
||||||
|
content: content.to_string(),
|
||||||
|
model: model.map(|m| m.to_string()),
|
||||||
|
timestamp: now.to_rfc3339(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn save(&self, base_dir: &Path) -> Result<(), io::Error> {
|
||||||
|
let sessions_dir = base_dir.join("sessions");
|
||||||
|
fs::create_dir_all(&sessions_dir)?;
|
||||||
|
|
||||||
|
let file_path = sessions_dir.join(format!("{}.json", self.id));
|
||||||
|
let json = serde_json::to_string_pretty(self)
|
||||||
|
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
|
||||||
|
fs::write(file_path, json)?;
|
||||||
|
|
||||||
|
// Also save latest symlink/copy for convenient continuation
|
||||||
|
let latest_path = sessions_dir.join("latest.json");
|
||||||
|
let _ = fs::write(
|
||||||
|
latest_path,
|
||||||
|
serde_json::to_string_pretty(self).unwrap_or_default(),
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn load(base_dir: &Path, id: &str) -> Result<Self, io::Error> {
|
||||||
|
let sessions_dir = base_dir.join("sessions");
|
||||||
|
let file_path = if id == "latest" {
|
||||||
|
sessions_dir.join("latest.json")
|
||||||
|
} else {
|
||||||
|
sessions_dir.join(format!("{id}.json"))
|
||||||
|
};
|
||||||
|
|
||||||
|
let content = fs::read_to_string(file_path)?;
|
||||||
|
let session: Session = serde_json::from_str(&content)
|
||||||
|
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
|
||||||
|
Ok(session)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn list(base_dir: &Path) -> Result<Vec<SessionSummary>, io::Error> {
|
||||||
|
let sessions_dir = base_dir.join("sessions");
|
||||||
|
if !sessions_dir.exists() {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut list = Vec::new();
|
||||||
|
for entry in fs::read_dir(sessions_dir)? {
|
||||||
|
let entry = entry?;
|
||||||
|
let path = entry.path();
|
||||||
|
if path.extension().and_then(|e| e.to_str()) == Some("json") {
|
||||||
|
let filename = path
|
||||||
|
.file_stem()
|
||||||
|
.and_then(|s| s.to_str())
|
||||||
|
.unwrap_or_default();
|
||||||
|
if filename == "latest" {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if let Ok(content) = fs::read_to_string(&path) {
|
||||||
|
if let Ok(session) = serde_json::from_str::<Session>(&content) {
|
||||||
|
list.push(SessionSummary {
|
||||||
|
id: session.id,
|
||||||
|
title: session.title,
|
||||||
|
persona: session.persona,
|
||||||
|
updated_at: session.updated_at,
|
||||||
|
message_count: session.messages.len(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
list.sort_by(|a, b| b.updated_at.cmp(&a.updated_at));
|
||||||
|
Ok(list)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn delete(base_dir: &Path, id: &str) -> Result<(), io::Error> {
|
||||||
|
let file_path = base_dir.join("sessions").join(format!("{id}.json"));
|
||||||
|
if file_path.exists() {
|
||||||
|
fs::remove_file(file_path)?;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use tempfile::tempdir;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_session_lifecycle_save_and_load() {
|
||||||
|
let dir = tempdir().unwrap();
|
||||||
|
let mut session = Session::new("coder");
|
||||||
|
session.add_message("you", "Write an LRU cache", None);
|
||||||
|
session.add_message("coder", "Here is the code", Some("qwen"));
|
||||||
|
|
||||||
|
assert_eq!(session.title, "Write an LRU cache");
|
||||||
|
session.save(dir.path()).expect("save succeeds");
|
||||||
|
|
||||||
|
let loaded = Session::load(dir.path(), &session.id).expect("load succeeds");
|
||||||
|
assert_eq!(loaded.messages.len(), 2);
|
||||||
|
assert_eq!(loaded.persona, "coder");
|
||||||
|
|
||||||
|
let summaries = Session::list(dir.path()).expect("list succeeds");
|
||||||
|
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");
|
||||||
|
}
|
||||||
|
}
|
||||||
182
src/tool.rs
Normal file
182
src/tool.rs
Normal file
|
|
@ -0,0 +1,182 @@
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::path::Path;
|
||||||
|
use std::time::Duration;
|
||||||
|
use tokio::process::Command;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ToolOutput {
|
||||||
|
pub command: String,
|
||||||
|
pub exit_code: i32,
|
||||||
|
pub stdout: String,
|
||||||
|
pub stderr: String,
|
||||||
|
pub duration_ms: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ToolOutput {
|
||||||
|
pub fn is_success(&self) -> bool {
|
||||||
|
self.exit_code == 0
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn format_for_model(&self) -> String {
|
||||||
|
let mut out = format!(
|
||||||
|
"[Command: `{}` exit code: {}]\n",
|
||||||
|
self.command, self.exit_code
|
||||||
|
);
|
||||||
|
if !self.stdout.is_empty() {
|
||||||
|
out.push_str("stdout:\n");
|
||||||
|
out.push_str(&self.stdout);
|
||||||
|
if !self.stdout.ends_with('\n') {
|
||||||
|
out.push('\n');
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !self.stderr.is_empty() {
|
||||||
|
out.push_str("stderr:\n");
|
||||||
|
out.push_str(&self.stderr);
|
||||||
|
if !self.stderr.ends_with('\n') {
|
||||||
|
out.push('\n');
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if self.stdout.is_empty() && self.stderr.is_empty() {
|
||||||
|
out.push_str("(No output produced)\n");
|
||||||
|
}
|
||||||
|
out
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ToolEngine;
|
||||||
|
|
||||||
|
impl ToolEngine {
|
||||||
|
/// Executes a bash command locally with timeout and output capture
|
||||||
|
pub async fn execute_bash(
|
||||||
|
command: &str,
|
||||||
|
cwd: Option<&Path>,
|
||||||
|
timeout_duration: Duration,
|
||||||
|
) -> Result<ToolOutput, String> {
|
||||||
|
let start = std::time::Instant::now();
|
||||||
|
|
||||||
|
let mut cmd = Command::new("bash");
|
||||||
|
cmd.arg("-c").arg(command);
|
||||||
|
|
||||||
|
if let Some(dir) = cwd {
|
||||||
|
cmd.current_dir(dir);
|
||||||
|
}
|
||||||
|
|
||||||
|
let child = cmd.output();
|
||||||
|
|
||||||
|
match tokio::time::timeout(timeout_duration, child).await {
|
||||||
|
Ok(Ok(output)) => {
|
||||||
|
let duration_ms = start.elapsed().as_millis() as u64;
|
||||||
|
let exit_code = output.status.code().unwrap_or(-1);
|
||||||
|
|
||||||
|
let stdout_raw = String::from_utf8_lossy(&output.stdout);
|
||||||
|
let stderr_raw = String::from_utf8_lossy(&output.stderr);
|
||||||
|
|
||||||
|
// Truncate to maximum 32KB to protect token context
|
||||||
|
let stdout = truncate_output(&stdout_raw, 32 * 1024);
|
||||||
|
let stderr = truncate_output(&stderr_raw, 16 * 1024);
|
||||||
|
|
||||||
|
Ok(ToolOutput {
|
||||||
|
command: command.to_string(),
|
||||||
|
exit_code,
|
||||||
|
stdout,
|
||||||
|
stderr,
|
||||||
|
duration_ms,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
Ok(Err(e)) => Err(format!("Process spawn error: {e}")),
|
||||||
|
Err(_) => Err(format!(
|
||||||
|
"Command timed out after {}s",
|
||||||
|
timeout_duration.as_secs()
|
||||||
|
)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Extracts any command execution requests formatted as:
|
||||||
|
/// `<run>cmd</run>`, `<bash>cmd</bash>`, or markdown ```bash blocks tagged for execution
|
||||||
|
pub fn extract_command_calls(text: &str) -> Vec<String> {
|
||||||
|
let mut commands = Vec::new();
|
||||||
|
|
||||||
|
// 1. Check for <run>...</run> tags
|
||||||
|
let mut rest = text;
|
||||||
|
while let Some(start_idx) = rest.find("<run>") {
|
||||||
|
let after_start = &rest[start_idx + 5..];
|
||||||
|
if let Some(end_idx) = after_start.find("</run>") {
|
||||||
|
let cmd = after_start[..end_idx].trim();
|
||||||
|
if !cmd.is_empty() {
|
||||||
|
commands.push(cmd.to_string());
|
||||||
|
}
|
||||||
|
rest = &after_start[end_idx + 6..];
|
||||||
|
} else {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. Check for <bash>...</bash> tags
|
||||||
|
let mut rest_bash = text;
|
||||||
|
while let Some(start_idx) = rest_bash.find("<bash>") {
|
||||||
|
let after_start = &rest_bash[start_idx + 6..];
|
||||||
|
if let Some(end_idx) = after_start.find("</bash>") {
|
||||||
|
let cmd = after_start[..end_idx].trim();
|
||||||
|
if !cmd.is_empty() && !commands.contains(&cmd.to_string()) {
|
||||||
|
commands.push(cmd.to_string());
|
||||||
|
}
|
||||||
|
rest_bash = &after_start[end_idx + 7..];
|
||||||
|
} else {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
commands
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Tool instructions appended to persona system prompts
|
||||||
|
pub fn system_tool_prompt() -> &'static str {
|
||||||
|
"\n\n[TOOL CAPABILITY: LOCAL BASH]\nYou can execute local shell commands on the user's system when needed (e.g. running tests, reading files, inspecting directories, running git/cargo/python/npm). To run a command, output it enclosed in <run>...</run> tags, for example:\n<run>cargo check</run>\nor\n<run>git status</run>\nThe command output will be executed and returned to you."
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn truncate_output(s: &str, max_bytes: usize) -> String {
|
||||||
|
if s.len() <= max_bytes {
|
||||||
|
s.to_string()
|
||||||
|
} else {
|
||||||
|
// 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])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_extract_command_calls() {
|
||||||
|
let text = "Let me check the repository status:\n<run>git status</run>\nand also list files:\n<bash>ls -la</bash>";
|
||||||
|
let calls = ToolEngine::extract_command_calls(text);
|
||||||
|
assert_eq!(calls, vec!["git status", "ls -la"]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_execute_echo_command() {
|
||||||
|
let res =
|
||||||
|
ToolEngine::execute_bash("echo 'hello from brokecrew'", None, Duration::from_secs(5))
|
||||||
|
.await
|
||||||
|
.expect("execution succeeds");
|
||||||
|
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");
|
||||||
|
}
|
||||||
|
}
|
||||||
670
src/tui/app.rs
670
src/tui/app.rs
|
|
@ -1,5 +1,8 @@
|
||||||
use brokecrew::persona::Tier;
|
use brokecrew::persona::Tier;
|
||||||
use brokecrew::pipeline::PipelineConfig;
|
use brokecrew::pipeline::PipelineConfig;
|
||||||
|
use brokecrew::session::{Session, SessionSummary};
|
||||||
|
use std::path::PathBuf;
|
||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
pub enum MessageSender {
|
pub enum MessageSender {
|
||||||
|
|
@ -9,12 +12,10 @@ pub enum MessageSender {
|
||||||
model: String,
|
model: String,
|
||||||
tier: Tier,
|
tier: Tier,
|
||||||
},
|
},
|
||||||
PipelineStep {
|
Tool {
|
||||||
step_index: usize,
|
command: String,
|
||||||
total_steps: usize,
|
exit_code: i32,
|
||||||
persona: String,
|
duration_ms: u64,
|
||||||
model: String,
|
|
||||||
tier: Tier,
|
|
||||||
},
|
},
|
||||||
System,
|
System,
|
||||||
}
|
}
|
||||||
|
|
@ -31,14 +32,195 @@ pub enum TuiMode {
|
||||||
Normal,
|
Normal,
|
||||||
PersonaSelect,
|
PersonaSelect,
|
||||||
PipelineSelect,
|
PipelineSelect,
|
||||||
|
SessionSelect,
|
||||||
|
ThemeSelect,
|
||||||
|
Settings,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct CommandSuggestion {
|
||||||
|
pub name: &'static str,
|
||||||
|
pub description: &'static str,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub const ALL_COMMANDS: &[CommandSuggestion] = &[
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/settings",
|
||||||
|
description: "Open settings and feature menu",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/provider",
|
||||||
|
description: "View or change LLM provider base URL",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/key",
|
||||||
|
description: "View or update OpenRouter API key",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/model",
|
||||||
|
description: "View or set Stage B router model",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/models",
|
||||||
|
description: "Search live cached models catalog",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/session",
|
||||||
|
description: "Resume or manage saved sessions",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/theme",
|
||||||
|
description: "Select UI color theme",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/new",
|
||||||
|
description: "Start a fresh session",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/usage",
|
||||||
|
description: "View rate limit and budget usage counters",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/stats",
|
||||||
|
description: "Detailed usage and token analytics",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/export",
|
||||||
|
description: "Export current chat to Markdown or JSON",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/config",
|
||||||
|
description: "View active system configuration",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/stealth",
|
||||||
|
description: "Toggle stealth candidate models",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/escalate",
|
||||||
|
description: "Toggle complexity escalation",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/persona",
|
||||||
|
description: "Select active persona or pass name",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/p",
|
||||||
|
description: "Quick switch persona (e.g. /p coder, /p next)",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/pipeline",
|
||||||
|
description: "Select active pipeline or pass name",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/clear",
|
||||||
|
description: "Clear current conversation screen",
|
||||||
|
},
|
||||||
|
CommandSuggestion {
|
||||||
|
name: "/help",
|
||||||
|
description: "Show command and shortcut help",
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
pub const THEMES: &[(&str, &str)] = &[
|
||||||
|
(
|
||||||
|
"opencode",
|
||||||
|
"Default OpenCode Dark - Deep slate, cyan accent, forest tools",
|
||||||
|
),
|
||||||
|
("warm", "Warm Amber - Dark charcoal, golden amber accents"),
|
||||||
|
("forest", "Forest Mint - Deep pine, emerald & mint accents"),
|
||||||
|
(
|
||||||
|
"cyberpunk",
|
||||||
|
"Cyberpunk - Midnight, neon pink & electric blue",
|
||||||
|
),
|
||||||
|
("mono", "Monochrome - High contrast grayscale minimal"),
|
||||||
|
(
|
||||||
|
"catppuccin",
|
||||||
|
"Catppuccin Mocha - Soothing lavender, peach & velvet navy",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"tokyo_night",
|
||||||
|
"Tokyo Night - Midnight indigo, vibrant neon blue & cyan",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"nord",
|
||||||
|
"Nord - Polar night arctic frost, glacial blue & snow",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"dracula",
|
||||||
|
"Dracula - Classic vampire dark slate, neon purple & green",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"solarized",
|
||||||
|
"Solarized Dark - Deep ocean teal, cyan & solar amber",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"gruvbox",
|
||||||
|
"Gruvbox Dark - Retro dark charcoal, warm bright orange",
|
||||||
|
),
|
||||||
|
];
|
||||||
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub struct PersonaItem {
|
pub struct PersonaItem {
|
||||||
pub name: String,
|
pub name: String,
|
||||||
pub description: String,
|
pub description: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct CandidateAttemptInfo {
|
||||||
|
pub model: String,
|
||||||
|
pub tier: Tier,
|
||||||
|
pub status: String,
|
||||||
|
pub is_error: bool,
|
||||||
|
pub is_success: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct PipelineStepTelemetry {
|
||||||
|
pub persona: String,
|
||||||
|
pub model: String,
|
||||||
|
pub tier: Tier,
|
||||||
|
pub is_active: bool,
|
||||||
|
pub is_completed: bool,
|
||||||
|
pub fell_back: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct WorkingsState {
|
||||||
|
pub active_persona: String,
|
||||||
|
pub active_pipeline: Option<String>,
|
||||||
|
pub route_stage: Option<String>,
|
||||||
|
pub route_details: Option<String>,
|
||||||
|
pub complexity_score: Option<u8>,
|
||||||
|
pub candidate_attempts: Vec<CandidateAttemptInfo>,
|
||||||
|
pub pipeline_steps: Vec<PipelineStepTelemetry>,
|
||||||
|
pub events: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for WorkingsState {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
active_persona: "auto".to_string(),
|
||||||
|
active_pipeline: None,
|
||||||
|
route_stage: None,
|
||||||
|
route_details: None,
|
||||||
|
complexity_score: None,
|
||||||
|
candidate_attempts: Vec::new(),
|
||||||
|
pipeline_steps: Vec::new(),
|
||||||
|
events: vec![
|
||||||
|
format!(
|
||||||
|
"{} brokecrew online",
|
||||||
|
chrono::Local::now().format("%H:%M:%S")
|
||||||
|
),
|
||||||
|
format!(
|
||||||
|
"{} catalog and budget initialized",
|
||||||
|
chrono::Local::now().format("%H:%M:%S")
|
||||||
|
),
|
||||||
|
],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub struct AppState {
|
pub struct AppState {
|
||||||
pub messages: Vec<ChatMessage>,
|
pub messages: Vec<ChatMessage>,
|
||||||
pub input: String,
|
pub input: String,
|
||||||
|
|
@ -50,7 +232,10 @@ pub struct AppState {
|
||||||
pub scroll_offset: u16,
|
pub scroll_offset: u16,
|
||||||
pub auto_scroll: bool,
|
pub auto_scroll: bool,
|
||||||
|
|
||||||
// Personas & selection (0 = [Auto-Route], 1.. = specific personas)
|
// Esc confirmation
|
||||||
|
pub esc_pressed_at: Option<Instant>,
|
||||||
|
|
||||||
|
// Personas & selection
|
||||||
pub persona_items: Vec<PersonaItem>,
|
pub persona_items: Vec<PersonaItem>,
|
||||||
pub selected_persona_idx: usize,
|
pub selected_persona_idx: usize,
|
||||||
|
|
||||||
|
|
@ -59,16 +244,36 @@ pub struct AppState {
|
||||||
pub selected_pipeline_idx: usize,
|
pub selected_pipeline_idx: usize,
|
||||||
pub active_pipeline_for_prompt: Option<String>,
|
pub active_pipeline_for_prompt: Option<String>,
|
||||||
|
|
||||||
|
// Sessions & persistence
|
||||||
|
pub active_session: Session,
|
||||||
|
pub session_summaries: Vec<SessionSummary>,
|
||||||
|
pub selected_session_idx: usize,
|
||||||
|
pub config_dir: PathBuf,
|
||||||
|
|
||||||
|
// Workings telemetry
|
||||||
|
pub workings: WorkingsState,
|
||||||
|
|
||||||
// Rate limits & budget
|
// Rate limits & budget
|
||||||
pub minute_requests: usize,
|
pub minute_requests: usize,
|
||||||
pub minute_limit: usize,
|
pub minute_limit: usize,
|
||||||
pub daily_requests: usize,
|
pub daily_requests: usize,
|
||||||
pub daily_limit: usize,
|
pub daily_limit: usize,
|
||||||
pub models_cached: usize,
|
pub models_cached: usize,
|
||||||
|
pub tokens_in: usize,
|
||||||
|
pub tokens_out: usize,
|
||||||
|
|
||||||
// Options
|
// Options
|
||||||
pub allow_stealth: bool,
|
pub allow_stealth: bool,
|
||||||
pub escalate: bool,
|
pub escalate: bool,
|
||||||
|
pub complexity_threshold: u8,
|
||||||
|
pub tool_execution_enabled: bool,
|
||||||
|
pub selected_setting_idx: usize,
|
||||||
|
|
||||||
|
// UX & Theme
|
||||||
|
pub expand_tool_outputs: bool,
|
||||||
|
pub current_theme: String,
|
||||||
|
pub selected_suggestion_idx: usize,
|
||||||
|
pub selected_theme_idx: usize,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl AppState {
|
impl AppState {
|
||||||
|
|
@ -77,12 +282,15 @@ impl AppState {
|
||||||
pipelines: Vec<PipelineConfig>,
|
pipelines: Vec<PipelineConfig>,
|
||||||
minute_limit: usize,
|
minute_limit: usize,
|
||||||
daily_limit: usize,
|
daily_limit: usize,
|
||||||
|
config_dir: PathBuf,
|
||||||
|
initial_theme: String,
|
||||||
|
complexity_threshold: u8,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
let mut persona_items = Vec::new();
|
let mut persona_items = Vec::new();
|
||||||
// Item 0 is always the smart Auto-Router
|
// Item 0 is always the smart Auto-Router
|
||||||
persona_items.push(PersonaItem {
|
persona_items.push(PersonaItem {
|
||||||
name: "Auto-Route".to_string(),
|
name: "auto".to_string(),
|
||||||
description: "Stage A keyword + Stage B LLM classifier".to_string(),
|
description: "keyword filter + LLM classifier".to_string(),
|
||||||
});
|
});
|
||||||
for (name, desc) in raw_personas {
|
for (name, desc) in raw_personas {
|
||||||
persona_items.push(PersonaItem {
|
persona_items.push(PersonaItem {
|
||||||
|
|
@ -91,42 +299,93 @@ impl AppState {
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
let session = Session::new("auto");
|
||||||
|
let session_summaries = Session::list(&config_dir).unwrap_or_default();
|
||||||
|
let selected_theme_idx = THEMES
|
||||||
|
.iter()
|
||||||
|
.position(|(name, _)| *name == initial_theme)
|
||||||
|
.unwrap_or(0);
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
messages: vec![ChatMessage {
|
messages: Vec::new(),
|
||||||
sender: MessageSender::System,
|
|
||||||
content: "⚡ BrokeCrew AI Team TUI initialized. Type your prompt below, or press [Tab] to pick a persona, [Ctrl+P] to run a pipeline.".to_string(),
|
|
||||||
fell_back: false,
|
|
||||||
}],
|
|
||||||
input: String::new(),
|
input: String::new(),
|
||||||
cursor_pos: 0,
|
cursor_pos: 0,
|
||||||
mode: TuiMode::Normal,
|
mode: TuiMode::Normal,
|
||||||
is_generating: false,
|
is_generating: false,
|
||||||
spinner_tick: 0,
|
spinner_tick: 0,
|
||||||
status_message: "Ready".to_string(),
|
status_message: "ready".to_string(),
|
||||||
scroll_offset: 0,
|
scroll_offset: 0,
|
||||||
auto_scroll: true,
|
auto_scroll: true,
|
||||||
|
esc_pressed_at: None,
|
||||||
persona_items,
|
persona_items,
|
||||||
selected_persona_idx: 0,
|
selected_persona_idx: 0,
|
||||||
pipelines,
|
pipelines,
|
||||||
selected_pipeline_idx: 0,
|
selected_pipeline_idx: 0,
|
||||||
active_pipeline_for_prompt: None,
|
active_pipeline_for_prompt: None,
|
||||||
|
active_session: session,
|
||||||
|
session_summaries,
|
||||||
|
selected_session_idx: 0,
|
||||||
|
config_dir,
|
||||||
|
workings: WorkingsState::default(),
|
||||||
minute_requests: 0,
|
minute_requests: 0,
|
||||||
minute_limit,
|
minute_limit,
|
||||||
daily_requests: 0,
|
daily_requests: 0,
|
||||||
daily_limit,
|
daily_limit,
|
||||||
models_cached: 0,
|
models_cached: 0,
|
||||||
|
tokens_in: 0,
|
||||||
|
tokens_out: 0,
|
||||||
allow_stealth: false,
|
allow_stealth: false,
|
||||||
escalate: false,
|
escalate: false,
|
||||||
|
complexity_threshold,
|
||||||
|
tool_execution_enabled: true,
|
||||||
|
selected_setting_idx: 0,
|
||||||
|
expand_tool_outputs: false,
|
||||||
|
current_theme: initial_theme,
|
||||||
|
selected_suggestion_idx: 0,
|
||||||
|
selected_theme_idx,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn format_tokens(tokens: usize) -> String {
|
||||||
|
if tokens < 1000 {
|
||||||
|
format!("{tokens}")
|
||||||
|
} else {
|
||||||
|
format!("{:.1}k", tokens as f64 / 1000.0)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn get_matching_suggestions(&self) -> Vec<&'static CommandSuggestion> {
|
||||||
|
if !self.input.starts_with('/') {
|
||||||
|
return Vec::new();
|
||||||
|
}
|
||||||
|
let query = self.input.to_lowercase();
|
||||||
|
ALL_COMMANDS
|
||||||
|
.iter()
|
||||||
|
.filter(|c| c.name.starts_with(&query))
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn is_start_screen(&self) -> bool {
|
||||||
|
self.messages.is_empty()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn cycle_theme(&mut self) {
|
||||||
|
let current_pos = THEMES
|
||||||
|
.iter()
|
||||||
|
.position(|(name, _)| *name == self.current_theme)
|
||||||
|
.unwrap_or(0);
|
||||||
|
let next_pos = (current_pos + 1) % THEMES.len();
|
||||||
|
self.current_theme = THEMES[next_pos].0.to_string();
|
||||||
|
self.selected_theme_idx = next_pos;
|
||||||
|
}
|
||||||
|
|
||||||
pub fn current_persona_name(&self) -> &str {
|
pub fn current_persona_name(&self) -> &str {
|
||||||
if self.selected_persona_idx == 0 {
|
if self.selected_persona_idx == 0 {
|
||||||
"Auto-Route"
|
"auto"
|
||||||
} else if let Some(item) = self.persona_items.get(self.selected_persona_idx) {
|
} else if let Some(item) = self.persona_items.get(self.selected_persona_idx) {
|
||||||
&item.name
|
&item.name
|
||||||
} else {
|
} else {
|
||||||
"Auto-Route"
|
"auto"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -134,7 +393,87 @@ impl AppState {
|
||||||
self.selected_persona_idx == 0
|
self.selected_persona_idx == 0
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn cycle_persona(&mut self) {
|
||||||
|
if !self.persona_items.is_empty() {
|
||||||
|
self.selected_persona_idx = (self.selected_persona_idx + 1) % self.persona_items.len();
|
||||||
|
self.active_pipeline_for_prompt = None;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn set_persona_by_name(&mut self, query: &str) -> Result<String, String> {
|
||||||
|
let q = query.trim().to_lowercase();
|
||||||
|
if q.is_empty() {
|
||||||
|
return Err("Persona name cannot be empty".to_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Exact match
|
||||||
|
if let Some(pos) = self
|
||||||
|
.persona_items
|
||||||
|
.iter()
|
||||||
|
.position(|p| p.name.to_lowercase() == q)
|
||||||
|
{
|
||||||
|
self.selected_persona_idx = pos;
|
||||||
|
self.active_pipeline_for_prompt = None;
|
||||||
|
return Ok(self.current_persona_name().to_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Prefix match
|
||||||
|
if let Some(pos) = self
|
||||||
|
.persona_items
|
||||||
|
.iter()
|
||||||
|
.position(|p| p.name.to_lowercase().starts_with(&q))
|
||||||
|
{
|
||||||
|
self.selected_persona_idx = pos;
|
||||||
|
self.active_pipeline_for_prompt = None;
|
||||||
|
return Ok(self.current_persona_name().to_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Substring match
|
||||||
|
if let Some(pos) = self
|
||||||
|
.persona_items
|
||||||
|
.iter()
|
||||||
|
.position(|p| p.name.to_lowercase().contains(&q))
|
||||||
|
{
|
||||||
|
self.selected_persona_idx = pos;
|
||||||
|
self.active_pipeline_for_prompt = None;
|
||||||
|
return Ok(self.current_persona_name().to_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
let names: Vec<String> = self.persona_items.iter().map(|p| p.name.clone()).collect();
|
||||||
|
Err(format!(
|
||||||
|
"No persona matching '{query}'. Available: {}",
|
||||||
|
names.join(", ")
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn log_working_event(&mut self, text: &str) {
|
||||||
|
let now = chrono::Local::now().format("%H:%M:%S").to_string();
|
||||||
|
self.workings.events.push(format!("{now} {text}"));
|
||||||
|
if self.workings.events.len() > 60 {
|
||||||
|
self.workings.events.remove(0);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn reset_turn_telemetry(&mut self, persona: &str, pipeline: Option<String>) {
|
||||||
|
self.workings.active_persona = persona.to_string();
|
||||||
|
self.workings.active_pipeline = pipeline;
|
||||||
|
self.workings.route_stage = None;
|
||||||
|
self.workings.route_details = None;
|
||||||
|
self.workings.complexity_score = None;
|
||||||
|
self.workings.candidate_attempts.clear();
|
||||||
|
self.workings.pipeline_steps.clear();
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn update_current_message_model(&mut self, model_id: &str) {
|
||||||
|
if let Some(msg) = self.messages.last_mut() {
|
||||||
|
if let MessageSender::Persona { model, .. } = &mut msg.sender {
|
||||||
|
*model = model_id.to_string();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub fn add_user_message(&mut self, text: String) {
|
pub fn add_user_message(&mut self, text: String) {
|
||||||
|
self.active_session.add_message("you", &text, None);
|
||||||
self.messages.push(ChatMessage {
|
self.messages.push(ChatMessage {
|
||||||
sender: MessageSender::User,
|
sender: MessageSender::User,
|
||||||
content: text,
|
content: text,
|
||||||
|
|
@ -160,23 +499,22 @@ impl AppState {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn start_pipeline_step_message(
|
pub fn add_tool_result_message(
|
||||||
&mut self,
|
&mut self,
|
||||||
step_index: usize,
|
command: String,
|
||||||
total_steps: usize,
|
exit_code: i32,
|
||||||
persona: &str,
|
output: String,
|
||||||
model: &str,
|
duration_ms: u64,
|
||||||
tier: Tier,
|
|
||||||
) {
|
) {
|
||||||
|
self.active_session
|
||||||
|
.add_message("tool", &format!("$ {}\n{}", command, output), None);
|
||||||
self.messages.push(ChatMessage {
|
self.messages.push(ChatMessage {
|
||||||
sender: MessageSender::PipelineStep {
|
sender: MessageSender::Tool {
|
||||||
step_index,
|
command,
|
||||||
total_steps,
|
exit_code,
|
||||||
persona: persona.to_string(),
|
duration_ms,
|
||||||
model: model.to_string(),
|
|
||||||
tier,
|
|
||||||
},
|
},
|
||||||
content: String::new(),
|
content: output,
|
||||||
fell_back: false,
|
fell_back: false,
|
||||||
});
|
});
|
||||||
if self.auto_scroll {
|
if self.auto_scroll {
|
||||||
|
|
@ -204,20 +542,171 @@ impl AppState {
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn save_current_session(&mut self) {
|
||||||
|
self.active_session.messages.clear();
|
||||||
|
// Sync active_session messages
|
||||||
|
for msg in &self.messages {
|
||||||
|
let role = match &msg.sender {
|
||||||
|
MessageSender::User => "you",
|
||||||
|
MessageSender::Persona { name, .. } => name.as_str(),
|
||||||
|
MessageSender::Tool { .. } => "tool",
|
||||||
|
MessageSender::System => "system",
|
||||||
|
};
|
||||||
|
let model = match &msg.sender {
|
||||||
|
MessageSender::Persona { model, .. } => Some(model.as_str()),
|
||||||
|
_ => None,
|
||||||
|
};
|
||||||
|
self.active_session.add_message(role, &msg.content, model);
|
||||||
|
}
|
||||||
|
|
||||||
|
let _ = self.active_session.save(&self.config_dir);
|
||||||
|
self.session_summaries = Session::list(&self.config_dir).unwrap_or_default();
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn load_session_by_id(&mut self, session_id: &str) -> Result<(), String> {
|
||||||
|
let loaded = Session::load(&self.config_dir, session_id).map_err(|e| e.to_string())?;
|
||||||
|
self.messages.clear();
|
||||||
|
|
||||||
|
for m in &loaded.messages {
|
||||||
|
let sender = match m.sender.to_lowercase().as_str() {
|
||||||
|
"you" | "user" => MessageSender::User,
|
||||||
|
"system" => MessageSender::System,
|
||||||
|
"tool" => MessageSender::Tool {
|
||||||
|
command: "exec".to_string(),
|
||||||
|
exit_code: 0,
|
||||||
|
duration_ms: 0,
|
||||||
|
},
|
||||||
|
other => MessageSender::Persona {
|
||||||
|
name: other.to_string(),
|
||||||
|
model: m.model.clone().unwrap_or_else(|| "unknown".to_string()),
|
||||||
|
tier: Tier::Free,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
self.messages.push(ChatMessage {
|
||||||
|
sender,
|
||||||
|
content: m.content.clone(),
|
||||||
|
fell_back: false,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
self.tokens_in = self
|
||||||
|
.messages
|
||||||
|
.iter()
|
||||||
|
.filter(|m| matches!(m.sender, MessageSender::User))
|
||||||
|
.map(|m| brokecrew::estimate_tokens(&m.content))
|
||||||
|
.sum();
|
||||||
|
self.tokens_out = self
|
||||||
|
.messages
|
||||||
|
.iter()
|
||||||
|
.filter(|m| !matches!(m.sender, MessageSender::User | MessageSender::System))
|
||||||
|
.map(|m| brokecrew::estimate_tokens(&m.content))
|
||||||
|
.sum();
|
||||||
|
|
||||||
|
self.active_session = loaded;
|
||||||
|
self.scroll_offset = 0;
|
||||||
|
self.log_working_event(&format!("[session] loaded {}", session_id));
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn start_new_session(&mut self) {
|
||||||
|
let persona = self.current_persona_name().to_string();
|
||||||
|
self.active_session = Session::new(&persona);
|
||||||
|
self.messages.clear();
|
||||||
|
self.tokens_in = 0;
|
||||||
|
self.tokens_out = 0;
|
||||||
|
self.add_system_message("Started fresh session.".to_string());
|
||||||
|
self.scroll_offset = 0;
|
||||||
|
self.log_working_event("[session] started fresh session");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Exports active conversation to Markdown file in `~/.config/brokecrew/exports/`
|
||||||
|
pub fn export_active_session_md(&self) -> Result<PathBuf, std::io::Error> {
|
||||||
|
let export_dir = self.config_dir.join("exports");
|
||||||
|
std::fs::create_dir_all(&export_dir)?;
|
||||||
|
|
||||||
|
let filename = format!("{}.md", self.active_session.id);
|
||||||
|
let path = export_dir.join(filename);
|
||||||
|
|
||||||
|
let mut content = format!("# BrokeCrew Session: {}\n", self.active_session.title);
|
||||||
|
content.push_str(&format!("- **ID:** `{}`\n", self.active_session.id));
|
||||||
|
content.push_str(&format!(
|
||||||
|
"- **Created:** {}\n",
|
||||||
|
self.active_session.created_at
|
||||||
|
));
|
||||||
|
content.push_str(&format!(
|
||||||
|
"- **Updated:** {}\n\n---\n\n",
|
||||||
|
self.active_session.updated_at
|
||||||
|
));
|
||||||
|
|
||||||
|
for msg in &self.messages {
|
||||||
|
match &msg.sender {
|
||||||
|
MessageSender::User => {
|
||||||
|
content.push_str(&format!("### User\n\n{}\n\n", msg.content.trim()));
|
||||||
|
}
|
||||||
|
MessageSender::Persona { name, model, tier } => {
|
||||||
|
let tier_s = match tier {
|
||||||
|
Tier::Free => "free",
|
||||||
|
Tier::Paid => "paid",
|
||||||
|
};
|
||||||
|
content.push_str(&format!(
|
||||||
|
"### Assistant ({} • {} [{}])\n\n{}\n\n",
|
||||||
|
name,
|
||||||
|
model,
|
||||||
|
tier_s,
|
||||||
|
msg.content.trim()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
MessageSender::Tool {
|
||||||
|
command,
|
||||||
|
exit_code,
|
||||||
|
duration_ms,
|
||||||
|
} => {
|
||||||
|
let dur_s = (*duration_ms as f64) / 1000.0;
|
||||||
|
content.push_str(&format!(
|
||||||
|
"```bash\n$ {command}\n# exit {exit_code} ({dur_s:.1}s)\n{}\n```\n\n",
|
||||||
|
msg.content.trim()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
MessageSender::System => {
|
||||||
|
content.push_str(&format!("> *System: {}*\n\n", msg.content.trim()));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
std::fs::write(&path, content)?;
|
||||||
|
Ok(path)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Exports active conversation to JSON file in `~/.config/brokecrew/exports/`
|
||||||
|
pub fn export_active_session_json(&self) -> Result<PathBuf, std::io::Error> {
|
||||||
|
let export_dir = self.config_dir.join("exports");
|
||||||
|
std::fs::create_dir_all(&export_dir)?;
|
||||||
|
|
||||||
|
let filename = format!("{}.json", self.active_session.id);
|
||||||
|
let path = export_dir.join(filename);
|
||||||
|
|
||||||
|
let json = serde_json::to_string_pretty(&self.active_session)
|
||||||
|
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
|
||||||
|
std::fs::write(&path, json)?;
|
||||||
|
Ok(path)
|
||||||
|
}
|
||||||
|
|
||||||
pub fn clear_chat(&mut self) {
|
pub fn clear_chat(&mut self) {
|
||||||
self.messages.clear();
|
self.messages.clear();
|
||||||
self.add_system_message("Chat history cleared.".to_string());
|
self.add_system_message("Chat history cleared.".to_string());
|
||||||
self.scroll_offset = 0;
|
self.scroll_offset = 0;
|
||||||
|
self.log_working_event("[chat] history cleared");
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn insert_char(&mut self, c: char) {
|
pub fn insert_char(&mut self, c: char) {
|
||||||
|
self.esc_pressed_at = None;
|
||||||
self.input.insert(self.cursor_pos, c);
|
self.input.insert(self.cursor_pos, c);
|
||||||
self.cursor_pos += c.len_utf8();
|
self.cursor_pos += c.len_utf8();
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn delete_char_before_cursor(&mut self) {
|
pub fn delete_char_before_cursor(&mut self) {
|
||||||
|
self.esc_pressed_at = None;
|
||||||
if self.cursor_pos > 0 {
|
if self.cursor_pos > 0 {
|
||||||
// Find char boundary
|
|
||||||
let mut prev = self.cursor_pos - 1;
|
let mut prev = self.cursor_pos - 1;
|
||||||
while !self.input.is_char_boundary(prev) && prev > 0 {
|
while !self.input.is_char_boundary(prev) && prev > 0 {
|
||||||
prev -= 1;
|
prev -= 1;
|
||||||
|
|
@ -228,12 +717,14 @@ impl AppState {
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn delete_char_under_cursor(&mut self) {
|
pub fn delete_char_under_cursor(&mut self) {
|
||||||
|
self.esc_pressed_at = None;
|
||||||
if self.cursor_pos < self.input.len() {
|
if self.cursor_pos < self.input.len() {
|
||||||
self.input.remove(self.cursor_pos);
|
self.input.remove(self.cursor_pos);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn move_cursor_left(&mut self) {
|
pub fn move_cursor_left(&mut self) {
|
||||||
|
self.esc_pressed_at = None;
|
||||||
if self.cursor_pos > 0 {
|
if self.cursor_pos > 0 {
|
||||||
let mut prev = self.cursor_pos - 1;
|
let mut prev = self.cursor_pos - 1;
|
||||||
while !self.input.is_char_boundary(prev) && prev > 0 {
|
while !self.input.is_char_boundary(prev) && prev > 0 {
|
||||||
|
|
@ -244,6 +735,7 @@ impl AppState {
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn move_cursor_right(&mut self) {
|
pub fn move_cursor_right(&mut self) {
|
||||||
|
self.esc_pressed_at = None;
|
||||||
if self.cursor_pos < self.input.len() {
|
if self.cursor_pos < self.input.len() {
|
||||||
let mut next = self.cursor_pos + 1;
|
let mut next = self.cursor_pos + 1;
|
||||||
while !self.input.is_char_boundary(next) && next < self.input.len() {
|
while !self.input.is_char_boundary(next) && next < self.input.len() {
|
||||||
|
|
@ -254,10 +746,12 @@ impl AppState {
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn move_cursor_start(&mut self) {
|
pub fn move_cursor_start(&mut self) {
|
||||||
|
self.esc_pressed_at = None;
|
||||||
self.cursor_pos = 0;
|
self.cursor_pos = 0;
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn move_cursor_end(&mut self) {
|
pub fn move_cursor_end(&mut self) {
|
||||||
|
self.esc_pressed_at = None;
|
||||||
self.cursor_pos = self.input.len();
|
self.cursor_pos = self.input.len();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -273,3 +767,117 @@ impl AppState {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use std::path::PathBuf;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_app_start_screen_and_suggestions() {
|
||||||
|
let mut app = AppState::new(
|
||||||
|
vec![],
|
||||||
|
vec![],
|
||||||
|
20,
|
||||||
|
50,
|
||||||
|
PathBuf::from("/tmp/brokecrew_test"),
|
||||||
|
"opencode".to_string(),
|
||||||
|
7,
|
||||||
|
);
|
||||||
|
assert!(app.is_start_screen());
|
||||||
|
|
||||||
|
app.input = "/se".to_string();
|
||||||
|
let sugs = app.get_matching_suggestions();
|
||||||
|
assert!(!sugs.is_empty());
|
||||||
|
assert!(sugs.iter().any(|s| s.name == "/session"));
|
||||||
|
assert!(sugs.iter().any(|s| s.name == "/settings"));
|
||||||
|
|
||||||
|
app.input = "hello world".to_string();
|
||||||
|
let sugs2 = app.get_matching_suggestions();
|
||||||
|
assert!(sugs2.is_empty());
|
||||||
|
|
||||||
|
app.add_user_message("hello".to_string());
|
||||||
|
assert!(!app.is_start_screen());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_app_theme_cycling() {
|
||||||
|
let mut app = AppState::new(
|
||||||
|
vec![],
|
||||||
|
vec![],
|
||||||
|
20,
|
||||||
|
50,
|
||||||
|
PathBuf::from("/tmp/brokecrew_test"),
|
||||||
|
"opencode".to_string(),
|
||||||
|
7,
|
||||||
|
);
|
||||||
|
assert_eq!(app.current_theme, "opencode");
|
||||||
|
app.cycle_theme();
|
||||||
|
assert_eq!(app.current_theme, "warm");
|
||||||
|
app.cycle_theme();
|
||||||
|
assert_eq!(app.current_theme, "forest");
|
||||||
|
app.cycle_theme();
|
||||||
|
assert_eq!(app.current_theme, "cyberpunk");
|
||||||
|
app.cycle_theme();
|
||||||
|
assert_eq!(app.current_theme, "mono");
|
||||||
|
app.cycle_theme();
|
||||||
|
assert_eq!(app.current_theme, "catppuccin");
|
||||||
|
app.cycle_theme();
|
||||||
|
assert_eq!(app.current_theme, "tokyo_night");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_app_export() {
|
||||||
|
let temp_dir = tempfile::tempdir().unwrap();
|
||||||
|
let mut app = AppState::new(
|
||||||
|
vec![],
|
||||||
|
vec![],
|
||||||
|
20,
|
||||||
|
50,
|
||||||
|
temp_dir.path().to_path_buf(),
|
||||||
|
"opencode".to_string(),
|
||||||
|
7,
|
||||||
|
);
|
||||||
|
app.add_user_message("Hello from user".to_string());
|
||||||
|
app.start_assistant_message("coder", "qwen-coder", Tier::Free);
|
||||||
|
app.append_to_current_message("Hello from coder");
|
||||||
|
|
||||||
|
let md_path = app.export_active_session_md().expect("export md");
|
||||||
|
assert!(md_path.exists());
|
||||||
|
let md_content = std::fs::read_to_string(&md_path).unwrap();
|
||||||
|
assert!(md_content.contains("Hello from user"));
|
||||||
|
assert!(md_content.contains("Hello from coder"));
|
||||||
|
|
||||||
|
let json_path = app.export_active_session_json().expect("export json");
|
||||||
|
assert!(json_path.exists());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_persona_quick_switch() {
|
||||||
|
let personas = vec![
|
||||||
|
("coder".to_string(), "coding assistant".to_string()),
|
||||||
|
("reviewer".to_string(), "review assistant".to_string()),
|
||||||
|
];
|
||||||
|
let mut app = AppState::new(
|
||||||
|
personas,
|
||||||
|
vec![],
|
||||||
|
20,
|
||||||
|
50,
|
||||||
|
PathBuf::from("/tmp/brokecrew_test"),
|
||||||
|
"opencode".to_string(),
|
||||||
|
7,
|
||||||
|
);
|
||||||
|
assert_eq!(app.current_persona_name(), "auto");
|
||||||
|
|
||||||
|
assert_eq!(app.set_persona_by_name("c").unwrap(), "coder");
|
||||||
|
assert_eq!(app.current_persona_name(), "coder");
|
||||||
|
|
||||||
|
assert_eq!(app.set_persona_by_name("rev").unwrap(), "reviewer");
|
||||||
|
assert_eq!(app.current_persona_name(), "reviewer");
|
||||||
|
|
||||||
|
assert_eq!(app.set_persona_by_name("auto").unwrap(), "auto");
|
||||||
|
assert_eq!(app.current_persona_name(), "auto");
|
||||||
|
|
||||||
|
assert!(app.set_persona_by_name("nonexistent").is_err());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
||||||
1234
src/tui/mod.rs
1234
src/tui/mod.rs
File diff suppressed because it is too large
Load diff
1763
src/tui/ui.rs
1763
src/tui/ui.rs
File diff suppressed because it is too large
Load diff
|
|
@ -1,7 +1,7 @@
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use brokecrew::client::{ClientError, ModelClient};
|
use brokecrew::client::{ClientError, ModelClient};
|
||||||
use brokecrew::persona::{ModelCandidate, Persona, PersonaRegistry, Tier};
|
use brokecrew::persona::{ModelCandidate, Persona, PersonaRegistry, Tier};
|
||||||
use brokecrew::pipeline::PipelineExecutor;
|
use brokecrew::pipeline::{DispatchParams, PipelineExecutor};
|
||||||
use brokecrew::registry::ModelRegistry;
|
use brokecrew::registry::ModelRegistry;
|
||||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
@ -84,16 +84,16 @@ async fn test_candidate_cascade_on_503_service_unavailable() {
|
||||||
calls: calls.clone(),
|
calls: calls.clone(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let result = PipelineExecutor::dispatch_persona(
|
let result = PipelineExecutor::dispatch_persona(DispatchParams {
|
||||||
"resilient-coder",
|
persona_name: "resilient-coder",
|
||||||
&persona_reg,
|
persona_registry: &persona_reg,
|
||||||
&model_reg,
|
model_registry: &model_reg,
|
||||||
&client,
|
client: &client,
|
||||||
"write rust code",
|
prompt: "write rust code",
|
||||||
false,
|
allow_stealth: false,
|
||||||
false,
|
require_paid: false,
|
||||||
None,
|
chunk_sender: None,
|
||||||
)
|
})
|
||||||
.await
|
.await
|
||||||
.expect("fallback to candidate-2 should succeed");
|
.expect("fallback to candidate-2 should succeed");
|
||||||
|
|
||||||
|
|
@ -121,16 +121,16 @@ async fn test_candidate_cascade_on_429_rate_limited() {
|
||||||
calls: calls.clone(),
|
calls: calls.clone(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let result = PipelineExecutor::dispatch_persona(
|
let result = PipelineExecutor::dispatch_persona(DispatchParams {
|
||||||
"resilient-coder",
|
persona_name: "resilient-coder",
|
||||||
&persona_reg,
|
persona_registry: &persona_reg,
|
||||||
&model_reg,
|
model_registry: &model_reg,
|
||||||
&client,
|
client: &client,
|
||||||
"write python code",
|
prompt: "write python code",
|
||||||
false,
|
allow_stealth: false,
|
||||||
false,
|
require_paid: false,
|
||||||
None,
|
chunk_sender: None,
|
||||||
)
|
})
|
||||||
.await
|
.await
|
||||||
.expect("fallback to candidate-2 should succeed");
|
.expect("fallback to candidate-2 should succeed");
|
||||||
|
|
||||||
|
|
@ -151,16 +151,16 @@ async fn test_fail_fast_on_401_unauthorized() {
|
||||||
calls: calls.clone(),
|
calls: calls.clone(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let err = PipelineExecutor::dispatch_persona(
|
let err = PipelineExecutor::dispatch_persona(DispatchParams {
|
||||||
"resilient-coder",
|
persona_name: "resilient-coder",
|
||||||
&persona_reg,
|
persona_registry: &persona_reg,
|
||||||
&model_reg,
|
model_registry: &model_reg,
|
||||||
&client,
|
client: &client,
|
||||||
"test query",
|
prompt: "test query",
|
||||||
false,
|
allow_stealth: false,
|
||||||
false,
|
require_paid: false,
|
||||||
None,
|
chunk_sender: None,
|
||||||
)
|
})
|
||||||
.await
|
.await
|
||||||
.unwrap_err();
|
.unwrap_err();
|
||||||
|
|
||||||
|
|
@ -208,16 +208,16 @@ async fn test_total_cascade_failure_reports_all_attempts() {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
let err = PipelineExecutor::dispatch_persona(
|
let err = PipelineExecutor::dispatch_persona(DispatchParams {
|
||||||
"resilient-coder",
|
persona_name: "resilient-coder",
|
||||||
&persona_reg,
|
persona_registry: &persona_reg,
|
||||||
&model_reg,
|
model_registry: &model_reg,
|
||||||
&AllFailingClient,
|
client: &AllFailingClient,
|
||||||
"test query",
|
prompt: "test query",
|
||||||
false,
|
allow_stealth: false,
|
||||||
false,
|
require_paid: false,
|
||||||
None,
|
chunk_sender: None,
|
||||||
)
|
})
|
||||||
.await
|
.await
|
||||||
.unwrap_err();
|
.unwrap_err();
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue