brokecrew/tests/fallback_tests.rs

use async_trait::async_trait;
use brokecrew::client::{ClientError, ModelClient};
use brokecrew::persona::{ModelCandidate, Persona, PersonaRegistry, Tier};
use brokecrew::pipeline::{DispatchParams, PipelineExecutor};
use brokecrew::registry::ModelRegistry;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tokio::sync::mpsc;

struct FaultInjectingClient {
    candidate1_error: Option<ClientError>,
    calls: Arc<AtomicUsize>,
}

#[async_trait]
impl ModelClient for FaultInjectingClient {
    async fn complete(&self, model: &str, system: &str, user: &str) -> Result<String, ClientError> {
        self.calls.fetch_add(1, Ordering::SeqCst);
        if model == "candidate-1" {
            if let Some(ref err) = self.candidate1_error {
                return match err {
                    ClientError::RateLimited(s) => Err(ClientError::RateLimited(*s)),
                    ClientError::ServerError { status, message } => Err(ClientError::ServerError {
                        status: *status,
                        message: message.clone(),
                    }),
                    ClientError::Unauthorized => Err(ClientError::Unauthorized),
                    ClientError::PaymentRequired => Err(ClientError::PaymentRequired),
                    _ => Err(ClientError::ModelUnavailable("simulated".into())),
                };
            }
        }

        Ok(format!(
            "Success from {model} for prompt '{user}' (system: {system})"
        ))
    }

    async fn complete_stream(
        &self,
        model: &str,
        system: &str,
        user: &str,
        chunk_sender: mpsc::Sender<String>,
    ) -> Result<String, ClientError> {
        let reply = self.complete(model, system, user).await?;
        let _ = chunk_sender.send(reply.clone()).await;
        Ok(reply)
    }
}

fn create_resilience_persona() -> Persona {
    Persona {
        name: "resilient-coder".into(),
        system_prompt: "You are coder".into(),
        trigger_words: vec!["code".into()],
        candidates: vec![
            ModelCandidate {
                id: "candidate-1".into(),
                tier: Tier::Free,
                is_stealth: false,
            },
            ModelCandidate {
                id: "candidate-2".into(),
                tier: Tier::Free,
                is_stealth: false,
            },
        ],
    }
}

#[tokio::test]
async fn test_candidate_cascade_on_503_service_unavailable() {
    let persona = create_resilience_persona();
    let persona_reg = PersonaRegistry::new(vec![persona]);
    let model_reg = ModelRegistry::new();

    let calls = Arc::new(AtomicUsize::new(0));
    let client = FaultInjectingClient {
        candidate1_error: Some(ClientError::ServerError {
            status: 503,
            message: "Overloaded".into(),
        }),
        calls: calls.clone(),
    };

    let result = PipelineExecutor::dispatch_persona(DispatchParams {
        persona_name: "resilient-coder",
        persona_registry: &persona_reg,
        model_registry: &model_reg,
        client: &client,
        prompt: "write rust code",
        allow_stealth: false,
        require_paid: false,
        chunk_sender: None,
    })
    .await
    .expect("fallback to candidate-2 should succeed");

    assert_eq!(result.persona, "resilient-coder");
    assert_eq!(result.model_used, "candidate-2");
    assert!(result.fell_back, "Result must indicate fallback occurred");
    assert_eq!(result.attempts.len(), 1);
    assert_eq!(result.attempts[0].model_id, "candidate-1");
    assert_eq!(
        calls.load(Ordering::SeqCst),
        2,
        "Both candidates must have been called"
    );
}

#[tokio::test]
async fn test_candidate_cascade_on_429_rate_limited() {
    let persona = create_resilience_persona();
    let persona_reg = PersonaRegistry::new(vec![persona]);
    let model_reg = ModelRegistry::new();

    let calls = Arc::new(AtomicUsize::new(0));
    let client = FaultInjectingClient {
        candidate1_error: Some(ClientError::RateLimited(15)),
        calls: calls.clone(),
    };

    let result = PipelineExecutor::dispatch_persona(DispatchParams {
        persona_name: "resilient-coder",
        persona_registry: &persona_reg,
        model_registry: &model_reg,
        client: &client,
        prompt: "write python code",
        allow_stealth: false,
        require_paid: false,
        chunk_sender: None,
    })
    .await
    .expect("fallback to candidate-2 should succeed");

    assert_eq!(result.model_used, "candidate-2");
    assert!(result.fell_back);
    assert_eq!(calls.load(Ordering::SeqCst), 2);
}

#[tokio::test]
async fn test_fail_fast_on_401_unauthorized() {
    let persona = create_resilience_persona();
    let persona_reg = PersonaRegistry::new(vec![persona]);
    let model_reg = ModelRegistry::new();

    let calls = Arc::new(AtomicUsize::new(0));
    let client = FaultInjectingClient {
        candidate1_error: Some(ClientError::Unauthorized),
        calls: calls.clone(),
    };

    let err = PipelineExecutor::dispatch_persona(DispatchParams {
        persona_name: "resilient-coder",
        persona_registry: &persona_reg,
        model_registry: &model_reg,
        client: &client,
        prompt: "test query",
        allow_stealth: false,
        require_paid: false,
        chunk_sender: None,
    })
    .await
    .unwrap_err();

    match err {
        ClientError::Unauthorized => (),
        _ => panic!("Expected Unauthorized fail-fast, got {:?}", err),
    }

    // Must NOT call candidate-2 when 401 Unauthorized occurs!
    assert_eq!(calls.load(Ordering::SeqCst), 1);
}

#[tokio::test]
async fn test_total_cascade_failure_reports_all_attempts() {
    let persona = create_resilience_persona();
    let persona_reg = PersonaRegistry::new(vec![persona]);
    let model_reg = ModelRegistry::new();

    struct AllFailingClient;
    #[async_trait]
    impl ModelClient for AllFailingClient {
        async fn complete(
            &self,
            model: &str,
            _system: &str,
            _user: &str,
        ) -> Result<String, ClientError> {
            if model == "candidate-1" {
                Err(ClientError::ServerError {
                    status: 503,
                    message: "Model overloaded".into(),
                })
            } else {
                Err(ClientError::RateLimited(30))
            }
        }
        async fn complete_stream(
            &self,
            model: &str,
            system: &str,
            user: &str,
            _chunk_sender: mpsc::Sender<String>,
        ) -> Result<String, ClientError> {
            self.complete(model, system, user).await
        }
    }

    let err = PipelineExecutor::dispatch_persona(DispatchParams {
        persona_name: "resilient-coder",
        persona_registry: &persona_reg,
        model_registry: &model_reg,
        client: &AllFailingClient,
        prompt: "test query",
        allow_stealth: false,
        require_paid: false,
        chunk_sender: None,
    })
    .await
    .unwrap_err();

    let err_str = err.to_string();
    assert!(
        err_str.contains("candidate-1"),
        "Error must mention candidate-1: {err_str}"
    );
    assert!(
        err_str.contains("candidate-2"),
        "Error must mention candidate-2: {err_str}"
    );
}

#[tokio::test]
async fn test_stage_b_router_fails_fast_on_401() {
    use brokecrew::router::LlmRouter;

    struct AuthFailingClient;
    #[async_trait]
    impl ModelClient for AuthFailingClient {
        async fn complete(
            &self,
            _model: &str,
            _system: &str,
            _user: &str,
        ) -> Result<String, ClientError> {
            Err(ClientError::Unauthorized)
        }
        async fn complete_stream(
            &self,
            model: &str,
            system: &str,
            user: &str,
            _chunk_sender: mpsc::Sender<String>,
        ) -> Result<String, ClientError> {
            self.complete(model, system, user).await
        }
    }

    let err = LlmRouter::classify(
        &AuthFailingClient,
        "router-model",
        &["coder".into()],
        "general",
        "ambiguous prompt",
    )
    .await
    .unwrap_err();

    match err {
        ClientError::Unauthorized => (),
        _ => panic!("Expected Unauthorized from router, got {:?}", err),
    }
}