brokecrew/src/persona.rs

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Tier {
    Free,
    Paid,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ModelCandidate {
    pub id: String,
    pub tier: Tier,
    #[serde(default)]
    pub is_stealth: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Persona {
    pub name: String,
    pub system_prompt: String,
    #[serde(default)]
    pub trigger_words: Vec<String>,
    pub candidates: Vec<ModelCandidate>,
}

#[derive(Debug, Clone)]
pub struct ModelAttempt {
    pub model_id: String,
    pub error: Option<String>,
}

#[derive(Debug, Clone)]
pub struct DispatchResult {
    pub persona: String,
    pub model_used: String,
    pub tier_used: Tier,
    pub reply: String,
    pub fell_back: bool,
    pub attempts: Vec<ModelAttempt>,
}

#[derive(Debug, Clone, Default)]
pub struct PersonaRegistry {
    personas: Vec<Persona>,
}

impl PersonaRegistry {
    pub fn new(personas: Vec<Persona>) -> Self {
        Self { personas }
    }
}

impl From<Vec<Persona>> for PersonaRegistry {
    fn from(personas: Vec<Persona>) -> Self {
        Self::new(personas)
    }
}

impl PersonaRegistry {
    pub fn get(&self, name: &str) -> Option<&Persona> {
        self.personas.iter().find(|p| p.name == name)
    }

    pub fn list(&self) -> &[Persona] {
        &self.personas
    }

    /// Returns the ordered list of eligible model candidates for a persona.
    ///
    /// - If `allow_stealth` is false, excludes candidates where `is_stealth == true`.
    /// - If `require_paid` is true (e.g. complexity escalation), filters for `Tier::Paid` candidates.
    pub fn get_eligible_candidates(
        &self,
        persona_name: &str,
        allow_stealth: bool,
        require_paid: bool,
    ) -> Result<Vec<&ModelCandidate>, String> {
        let persona = self
            .get(persona_name)
            .ok_or_else(|| format!("Persona '{persona_name}' not found"))?;

        let eligible: Vec<&ModelCandidate> = persona
            .candidates
            .iter()
            .filter(|c| {
                if !allow_stealth && c.is_stealth {
                    return false;
                }
                if require_paid && c.tier != Tier::Paid {
                    return false;
                }
                true
            })
            .collect();

        if eligible.is_empty() {
            Err(format!(
                "No eligible model candidates for persona '{persona_name}' (stealth: {allow_stealth}, paid: {require_paid})"
            ))
        } else {
            Ok(eligible)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_candidate_filtering_and_fallback_order() {
        let candidates = vec![
            ModelCandidate {
                id: "stealth-model:free".into(),
                tier: Tier::Free,
                is_stealth: true,
            },
            ModelCandidate {
                id: "qwen-free".into(),
                tier: Tier::Free,
                is_stealth: false,
            },
            ModelCandidate {
                id: "gemini-paid".into(),
                tier: Tier::Paid,
                is_stealth: false,
            },
        ];

        let persona = Persona {
            name: "coder".into(),
            system_prompt: "You are a coder".into(),
            trigger_words: vec!["code".into()],
            candidates,
        };

        let registry = PersonaRegistry::new(vec![persona]);

        // When allow_stealth is false and require_paid is false
        let free_no_stealth = registry
            .get_eligible_candidates("coder", false, false)
            .expect("should return candidates");
        assert_eq!(free_no_stealth.len(), 2);
        assert_eq!(free_no_stealth[0].id, "qwen-free");
        assert_eq!(free_no_stealth[1].id, "gemini-paid");

        // When allow_stealth is true
        let with_stealth = registry
            .get_eligible_candidates("coder", true, false)
            .expect("should return candidates");
        assert_eq!(with_stealth.len(), 3);
        assert_eq!(with_stealth[0].id, "stealth-model:free");

        // When require_paid is true (complexity escalation)
        let paid_only = registry
            .get_eligible_candidates("coder", true, true)
            .expect("should return candidates");
        assert_eq!(paid_only.len(), 1);
        assert_eq!(paid_only[0].id, "gemini-paid");

        // Non-existent persona returns error
        assert!(registry
            .get_eligible_candidates("unknown", false, false)
            .is_err());
    }
}