use crate::db::MissionDB;
use crate::types::{Slot, SlotConfig};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use tracing::{debug, info};
pub struct SlotManager {
slots: Arc<RwLock<HashMap<String, Slot>>>,
db: Arc<MissionDB>,
}
impl SlotManager {
pub fn new(db: Arc<MissionDB>) -> Self {
Self {
slots: Arc::new(RwLock::new(HashMap::new())),
db,
}
}
pub fn load_slots(&self, configs: Vec<SlotConfig>) {
let mut slots = self.slots.write().unwrap();
for config in configs {
let saved_session_id = self.db.get_slot_session(&config.id).ok().flatten();
let slot = Slot {
config: config.clone(),
session_id: saved_session_id,
};
info!(slot_id = %config.id, role = %config.role, "Slot loaded");
slots.insert(config.id.clone(), slot);
}
}
pub fn get_all_slots(&self) -> Vec<Slot> {
let slots = self.slots.read().unwrap();
slots.values().cloned().collect()
}
pub fn get_slot(&self, slot_id: &str) -> Option<Slot> {
let slots = self.slots.read().unwrap();
slots.get(slot_id).cloned()
}
pub fn find_slot_by_role(&self, role: &str) -> Option<Slot> {
let slots = self.slots.read().unwrap();
slots.values().find(|s| s.config.role == role).cloned()
}
pub fn get_slots_by_role(&self, role: &str) -> Vec<Slot> {
let slots = self.slots.read().unwrap();
slots
.values()
.filter(|s| s.config.role == role)
.cloned()
.collect()
}
pub fn update_session(&self, slot_id: &str, session_id: &str) {
let mut slots = self.slots.write().unwrap();
if let Some(slot) = slots.get_mut(slot_id) {
slot.session_id = Some(session_id.to_string());
let _ = self.db.set_slot_session(slot_id, session_id);
debug!(slot_id = %slot_id, session_id = %session_id, "Session updated");
}
}
pub fn reset_session(&self, slot_id: &str) {
let mut slots = self.slots.write().unwrap();
if let Some(slot) = slots.get_mut(slot_id) {
slot.session_id = None;
self.db.clear_slot_session(slot_id);
info!(slot_id = %slot_id, "Session reset");
}
}
pub fn get_stats(&self) -> SlotStats {
let slots = self.slots.read().unwrap();
let mut by_role: HashMap<String, usize> = HashMap::new();
for slot in slots.values() {
*by_role.entry(slot.config.role.clone()).or_insert(0) += 1;
}
SlotStats {
total: slots.len(),
by_role,
}
}
}
#[derive(Debug, Clone)]
pub struct SlotStats {
pub total: usize,
pub by_role: HashMap<String, usize>,
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
fn create_test_db() -> Arc<MissionDB> {
let dir = tempdir().unwrap();
let db_path = dir.path().join("test.db");
Arc::new(MissionDB::open(db_path).unwrap())
}
#[test]
fn test_load_and_get_slots() {
let db = create_test_db();
let manager = SlotManager::new(db);
let configs = vec![
SlotConfig {
id: "slot-1".to_string(),
role: "worker".to_string(),
description: "Worker slot".to_string(),
cwd: None,
mcp_config: None,
auto_start: None,
},
SlotConfig {
id: "slot-2".to_string(),
role: "worker".to_string(),
description: "Another worker".to_string(),
cwd: None,
mcp_config: None,
auto_start: None,
},
SlotConfig {
id: "slot-3".to_string(),
role: "specialist".to_string(),
description: "Specialist slot".to_string(),
cwd: None,
mcp_config: None,
auto_start: None,
},
];
manager.load_slots(configs);
let all = manager.get_all_slots();
assert_eq!(all.len(), 3);
let slot = manager.get_slot("slot-1").unwrap();
assert_eq!(slot.config.role, "worker");
let workers = manager.get_slots_by_role("worker");
assert_eq!(workers.len(), 2);
let specialist = manager.find_slot_by_role("specialist").unwrap();
assert_eq!(specialist.config.id, "slot-3");
}
#[test]
fn test_session_management() {
let db = create_test_db();
let manager = SlotManager::new(db);
let configs = vec![SlotConfig {
id: "slot-1".to_string(),
role: "worker".to_string(),
description: "Worker slot".to_string(),
cwd: None,
mcp_config: None,
auto_start: None,
}];
manager.load_slots(configs);
let slot = manager.get_slot("slot-1").unwrap();
assert!(slot.session_id.is_none());
manager.update_session("slot-1", "session-abc");
let slot = manager.get_slot("slot-1").unwrap();
assert_eq!(slot.session_id, Some("session-abc".to_string()));
manager.reset_session("slot-1");
let slot = manager.get_slot("slot-1").unwrap();
assert!(slot.session_id.is_none());
}
#[test]
fn test_stats() {
let db = create_test_db();
let manager = SlotManager::new(db);
let configs = vec![
SlotConfig {
id: "slot-1".to_string(),
role: "worker".to_string(),
description: "Worker 1".to_string(),
cwd: None,
mcp_config: None,
auto_start: None,
},
SlotConfig {
id: "slot-2".to_string(),
role: "worker".to_string(),
description: "Worker 2".to_string(),
cwd: None,
mcp_config: None,
auto_start: None,
},
SlotConfig {
id: "slot-3".to_string(),
role: "specialist".to_string(),
description: "Specialist".to_string(),
cwd: None,
mcp_config: None,
auto_start: None,
},
];
manager.load_slots(configs);
let stats = manager.get_stats();
assert_eq!(stats.total, 3);
assert_eq!(stats.by_role.get("worker"), Some(&2));
assert_eq!(stats.by_role.get("specialist"), Some(&1));
}
}