use std::sync::Arc;
use anyhow::Result;
use parking_lot::RwLock;
use super::store::PersonaStore;
use super::{Persona, default_personas};
use crate::state_store::StateStore;
pub struct PersonaManager {
store: PersonaStore,
active_persona_id: RwLock<Option<String>>,
state_store: Option<Arc<StateStore>>,
#[allow(clippy::type_complexity)]
reseed_callback: RwLock<Option<Arc<dyn Fn(Option<String>) + Send + Sync>>>,
}
impl std::fmt::Debug for PersonaManager {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PersonaManager")
.field("active_persona_id", &self.active_persona_id.read())
.field("state_store", &self.state_store.is_some())
.field("reseed_callback", &self.reseed_callback.read().is_some())
.finish()
}
}
impl PersonaManager {
pub fn new() -> Self {
let store = PersonaStore::new();
let manager = Self {
store,
active_persona_id: RwLock::new(None),
state_store: None,
reseed_callback: RwLock::new(None),
};
manager.create_default_personas();
manager
}
pub fn with_defaults(personas: Vec<Persona>) -> Self {
let store = PersonaStore::new();
store.load_from_slice(&personas);
let this = Self {
store,
active_persona_id: RwLock::new(None),
state_store: None,
reseed_callback: RwLock::new(None),
};
if let Some(first) = this.store.list_enabled().into_iter().next() {
*this.active_persona_id.write() = Some(first.id);
}
this
}
pub fn with_state_store(mut self, store: Arc<StateStore>) -> Self {
self.state_store = Some(store);
self
}
pub fn set_reseed_callback(&self, cb: Option<Arc<dyn Fn(Option<String>) + Send + Sync>>) {
*self.reseed_callback.write() = cb;
}
pub fn get_active_persona(&self) -> Option<Persona> {
let active_id = self.active_persona_id.read().clone();
active_id.and_then(|id| self.store.get(&id))
}
pub fn active_system_prompt(&self) -> String {
self.get_active_persona()
.map(|p| p.system_prompt.clone())
.unwrap_or_else(|| {
"You are a helpful AI assistant that follows the Ouroboros methodology: \
specify before you build, evaluate before you ship."
.to_string()
})
}
pub fn create_default_personas(&self) {
let defaults = default_personas();
for persona in defaults {
if self.store.get(&persona.id).is_none() {
self.store.register(persona);
}
}
{
let mut active = self.active_persona_id.write();
if active.is_none() {
*active = Some("dev".to_string());
}
}
tracing::info!("Default personas initialized");
}
pub fn first_enabled(&self) -> Option<Persona> {
self.store.list_enabled().into_iter().next()
}
pub fn store(&self) -> &PersonaStore {
&self.store
}
pub fn active_persona_id(&self) -> Option<String> {
self.active_persona_id.read().clone()
}
pub fn apply_config(&self, cfg: &crate::config::PersonaConfig) {
if let Some(id) = self.active_persona_id()
&& self.store.get(&id).map(|p| p.enabled).unwrap_or(false)
{
return;
}
if let Some(id) = cfg.default_persona_id.as_ref()
&& self.store.get(id).map(|p| p.enabled).unwrap_or(false)
{
*self.active_persona_id.write() = Some(id.clone());
return;
}
if let Some(p) = self.store.list_enabled().into_iter().next() {
*self.active_persona_id.write() = Some(p.id);
}
}
pub async fn load_from_state_store(&self, store: &StateStore) -> Result<()> {
let snap = crate::persona::persistence::load_from_state_store(store)
.await?
.ok_or_else(|| anyhow::anyhow!("persona: no snapshot present"))?;
for p in &snap.personas {
self.store.register(p.clone());
}
if let Some(active) = snap.active_persona_id
&& snap.personas.iter().any(|p| p.id == active && p.enabled)
{
*self.active_persona_id.write() = Some(active);
}
Ok(())
}
pub async fn persist(&self) -> Result<()> {
let Some(ref store) = self.state_store else {
return Ok(());
};
let snapshot = crate::persona::persistence::PersonaSnapshot {
schema_version: 1,
active_persona_id: self.active_persona_id(),
personas: self.store.list_all(),
};
crate::persona::persistence::save_to_state_store(store, &snapshot).await
}
pub async fn set_active(&self, id: &str) -> Result<Option<String>> {
let persona = self
.store
.get(id)
.ok_or_else(|| anyhow::anyhow!("Persona '{id}' not found"))?;
if !persona.enabled {
anyhow::bail!("Persona '{id}' is disabled");
}
*self.active_persona_id.write() = Some(id.to_string());
tracing::info!(persona_id = %id, name = %persona.name, "Active persona set");
if let Err(e) = self.persist().await {
tracing::warn!(error = %e, "persona set_active: persist failed");
}
let prompt = persona.system_prompt.clone();
if let Some(ref cb) = *self.reseed_callback.read() {
cb(Some(prompt.clone()));
}
Ok(Some(prompt))
}
}
impl Default for PersonaManager {
fn default() -> Self {
Self::new()
}
}
impl Clone for PersonaManager {
fn clone(&self) -> Self {
let personas: Vec<Persona> = self.store.list_all();
let store = PersonaStore::new();
store.load_from_slice(&personas);
Self {
store,
active_persona_id: RwLock::new(self.active_persona_id.read().clone()),
state_store: self.state_store.clone(),
reseed_callback: RwLock::new(self.reseed_callback.read().clone()),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::PersonaConfig;
fn make_store() -> Arc<StateStore> {
let dir = tempfile::tempdir().unwrap();
Arc::new(StateStore::new(dir.keep()).unwrap())
}
#[tokio::test]
async fn test_load_from_state_store_round_trip() {
let store = make_store();
let pm = PersonaManager::new().with_state_store(store.clone());
let custom = Persona {
id: "custom-1".to_string(),
name: "Custom".to_string(),
role: "custom".to_string(),
description: "Custom test persona".to_string(),
system_prompt: "You are Custom.".to_string(),
enabled: true,
model: None,
personality_traits: vec![],
};
pm.store().register(custom);
pm.set_active("custom-1").await.unwrap();
pm.persist().await.unwrap();
let pm2 = PersonaManager::new().with_state_store(store.clone());
pm2.load_from_state_store(&store).await.unwrap();
assert!(pm2.store().get("custom-1").is_some());
assert_eq!(pm2.active_persona_id(), Some("custom-1".to_string()));
assert!(pm2.store().get("dev").is_some());
}
#[tokio::test]
async fn test_load_no_file_is_ok() {
let store = make_store();
let pm = PersonaManager::new().with_state_store(store.clone());
let result = pm.load_from_state_store(&store).await;
assert!(result.is_err());
assert!(pm.store().get("dev").is_some());
}
#[tokio::test]
async fn test_apply_config_default_persona_id() {
let pm = PersonaManager::new();
*pm.active_persona_id.write() = None;
let cfg = PersonaConfig {
default_persona_id: Some("review".to_string()),
};
pm.apply_config(&cfg);
assert_eq!(pm.active_persona_id(), Some("review".to_string()));
}
#[tokio::test]
async fn test_apply_config_falls_back_to_first_enabled() {
let pm = PersonaManager::new();
*pm.active_persona_id.write() = None;
let cfg = PersonaConfig {
default_persona_id: None,
};
pm.apply_config(&cfg);
assert_eq!(pm.active_persona_id(), Some("dev".to_string()));
}
#[tokio::test]
async fn test_apply_config_skips_disabled_default() {
let pm = PersonaManager::new();
*pm.active_persona_id.write() = None;
pm.store().set_enabled("review", false).unwrap();
let cfg = PersonaConfig {
default_persona_id: Some("review".to_string()),
};
pm.apply_config(&cfg);
assert_eq!(pm.active_persona_id(), Some("dev".to_string()));
}
#[tokio::test]
async fn test_apply_config_keeps_existing_active() {
let pm = PersonaManager::new();
pm.set_active("research").await.unwrap();
let cfg = PersonaConfig {
default_persona_id: Some("dev".to_string()),
};
pm.apply_config(&cfg);
assert_eq!(pm.active_persona_id(), Some("research".to_string()));
}
#[tokio::test]
async fn test_set_active_rejects_disabled() {
let pm = PersonaManager::new();
pm.store().set_enabled("review", false).unwrap();
let result = pm.set_active("review").await;
assert!(result.is_err());
assert_eq!(pm.active_persona_id(), Some("dev".to_string()));
}
#[tokio::test]
async fn test_set_active_rejects_unknown_id() {
let pm = PersonaManager::new();
let result = pm.set_active("nonexistent").await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_set_active_returns_system_prompt() {
let pm = PersonaManager::new();
let prompt = pm.set_active("review").await.unwrap();
assert!(prompt.is_some());
assert!(prompt.unwrap().contains("Review"));
}
#[tokio::test]
async fn test_set_active_fires_reseed_callback() {
let pm = PersonaManager::new();
let received = Arc::new(parking_lot::Mutex::new(None::<String>));
let received_cb = received.clone();
pm.set_reseed_callback(Some(Arc::new(move |prompt| {
*received_cb.lock() = prompt;
})));
pm.set_active("review").await.unwrap();
assert!(received.lock().as_ref().unwrap().contains("Review"));
}
#[tokio::test]
async fn test_set_active_no_callback_still_works() {
let pm = PersonaManager::new();
let result = pm.set_active("research").await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_set_active_persists_when_store_set() {
let store = make_store();
let pm = PersonaManager::new().with_state_store(store.clone());
pm.set_active("research").await.unwrap();
let pm2 = PersonaManager::new().with_state_store(store.clone());
pm2.load_from_state_store(&store).await.unwrap();
assert_eq!(pm2.active_persona_id(), Some("research".to_string()));
}
#[tokio::test]
async fn test_set_active_no_persist_without_store() {
let pm = PersonaManager::new();
pm.set_active("research").await.unwrap();
assert_eq!(pm.active_persona_id(), Some("research".to_string()));
}
#[tokio::test]
async fn test_persist_round_trip_preserves_active() {
let store = make_store();
let pm = PersonaManager::new().with_state_store(store.clone());
pm.set_active("research").await.unwrap();
pm.persist().await.unwrap();
let pm2 = PersonaManager::new().with_state_store(store.clone());
pm2.load_from_state_store(&store).await.unwrap();
assert_eq!(pm2.active_persona_id(), Some("research".to_string()));
}
#[tokio::test]
async fn test_clone_preserves_state_store_and_callback() {
let store = make_store();
let pm = PersonaManager::new().with_state_store(store.clone());
let received = Arc::new(parking_lot::Mutex::new(None::<String>));
let received_cb = received.clone();
pm.set_reseed_callback(Some(Arc::new(move |prompt| {
*received_cb.lock() = prompt;
})));
let cloned = pm.clone();
cloned.set_active("research").await.unwrap();
assert!(received.lock().is_some());
}
}