use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use anyhow::{bail, Context, Result};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
pub const DEFAULT_FLEET_CGROUP_ROOT: &str = "/sys/fs/cgroup/vetto-fleet";
pub const DEFAULT_MAX_AGENTS: usize = 64;
pub const DEFAULT_CPU_WEIGHT: u32 = 100;
pub const DEFAULT_MEMORY_LIMIT_BYTES: u64 = 2 * 1024 * 1024 * 1024; pub const DEFAULT_PIDS_MAX: u32 = 128;
pub const DEFAULT_BASE_PORT: u16 = 49201;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FleetConfig {
pub cgroup_root: PathBuf,
pub max_agents: usize,
pub cpu_weight: u32,
pub memory_limit_bytes: u64,
pub pids_max: u32,
pub base_port: u16,
pub ipc_isolation: bool,
}
impl Default for FleetConfig {
fn default() -> Self {
Self {
cgroup_root: PathBuf::from(DEFAULT_FLEET_CGROUP_ROOT),
max_agents: DEFAULT_MAX_AGENTS,
cpu_weight: DEFAULT_CPU_WEIGHT,
memory_limit_bytes: DEFAULT_MEMORY_LIMIT_BYTES,
pids_max: DEFAULT_PIDS_MAX,
base_port: DEFAULT_BASE_PORT,
ipc_isolation: true,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentWorkerScope {
pub worker_id: String,
pub agent_name: String,
pub scope_path: PathBuf,
pub cow_branch_name: String,
pub ephemeral_port: u16,
pub cpu_weight: u32,
pub memory_limit_bytes: u64,
pub pids_max: u32,
pub ipc_isolated: bool,
pub allocated_at: DateTime<Utc>,
}
#[derive(Debug, Clone)]
pub struct FleetManager {
config: FleetConfig,
active_workers: Arc<Mutex<BTreeMap<String, AgentWorkerScope>>>,
}
impl FleetManager {
pub fn new_default() -> Self {
Self::new(FleetConfig::default())
}
pub fn new(config: FleetConfig) -> Self {
Self {
config,
active_workers: Arc::new(Mutex::new(BTreeMap::new())),
}
}
pub fn config(&self) -> &FleetConfig {
&self.config
}
pub fn active_count(&self) -> usize {
self.active_workers.lock().map(|w| w.len()).unwrap_or(0)
}
pub fn allocate_worker(&self, agent_name: &str) -> Result<AgentWorkerScope> {
let mut workers = self
.active_workers
.lock()
.map_err(|_| anyhow::anyhow!("fleet manager lock poisoned"))?;
if workers.len() >= self.config.max_agents {
bail!(
"Fleet capacity exceeded: active agents ({}) >= max allowed ({})",
workers.len(),
self.config.max_agents
);
}
let mut slot_id = 1;
while workers.contains_key(&format!("agent-{:02}", slot_id)) {
slot_id += 1;
}
if slot_id > self.config.max_agents {
bail!(
"Fleet slot allocation exceeded max configured agents: slot {} > max {}",
slot_id,
self.config.max_agents
);
}
let worker_id = format!("agent-{:02}", slot_id);
let scope_path = self.config.cgroup_root.join(format!("{}.scope", worker_id));
let ephemeral_port = self
.config
.base_port
.checked_add((slot_id - 1) as u16)
.context("Ephemeral port pool overflow")?;
let scope = AgentWorkerScope {
worker_id: worker_id.clone(),
agent_name: agent_name.to_string(),
scope_path,
cow_branch_name: worker_id.clone(),
ephemeral_port,
cpu_weight: self.config.cpu_weight,
memory_limit_bytes: self.config.memory_limit_bytes,
pids_max: self.config.pids_max,
ipc_isolated: self.config.ipc_isolation,
allocated_at: Utc::now(),
};
workers.insert(worker_id, scope.clone());
Ok(scope)
}
pub fn release_worker(&self, worker_id: &str) -> Result<()> {
let mut workers = self
.active_workers
.lock()
.map_err(|_| anyhow::anyhow!("fleet manager lock poisoned"))?;
if workers.remove(worker_id).is_none() {
bail!("Worker scope '{}' not found in active fleet", worker_id);
}
Ok(())
}
pub fn get_worker(&self, worker_id: &str) -> Option<AgentWorkerScope> {
self.active_workers.lock().ok()?.get(worker_id).cloned()
}
pub fn all_workers(&self) -> Vec<AgentWorkerScope> {
self.active_workers
.lock()
.map(|w| w.values().cloned().collect())
.unwrap_or_default()
}
pub fn verify_isolation(&self, worker_a_id: &str, worker_b_id: &str) -> Result<()> {
if worker_a_id == worker_b_id {
bail!(
"Cannot verify isolation of a worker against itself ('{}')",
worker_a_id
);
}
let workers = self
.active_workers
.lock()
.map_err(|_| anyhow::anyhow!("fleet manager lock poisoned"))?;
let a = workers
.get(worker_a_id)
.ok_or_else(|| anyhow::anyhow!("Worker '{}' not found in active fleet", worker_a_id))?;
let b = workers
.get(worker_b_id)
.ok_or_else(|| anyhow::anyhow!("Worker '{}' not found in active fleet", worker_b_id))?;
if a.ephemeral_port == b.ephemeral_port {
bail!(
"Isolation breach: Workers '{}' and '{}' share ephemeral port {}",
worker_a_id,
worker_b_id,
a.ephemeral_port
);
}
if a.cow_branch_name == b.cow_branch_name {
bail!(
"Isolation breach: Workers '{}' and '{}' share CoW branch '{}'",
worker_a_id,
worker_b_id,
a.cow_branch_name
);
}
if a.scope_path == b.scope_path {
bail!(
"Isolation breach: Workers '{}' and '{}' share cgroup scope path '{:?}'",
worker_a_id,
worker_b_id,
a.scope_path
);
}
if !a.ipc_isolated || !b.ipc_isolated {
bail!(
"Isolation breach: IPC namespace isolation disabled for one or more workers ('{}': {}, '{}': {})",
worker_a_id,
a.ipc_isolated,
worker_b_id,
b.ipc_isolated
);
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_fleet_worker_allocation_and_isolation() {
let fleet = FleetManager::new_default();
let w1 = fleet.allocate_worker("claude-1").expect("allocate w1");
let w2 = fleet.allocate_worker("claude-2").expect("allocate w2");
assert_eq!(w1.worker_id, "agent-01");
assert_eq!(w2.worker_id, "agent-02");
assert_eq!(w1.ephemeral_port, DEFAULT_BASE_PORT);
assert_eq!(w2.ephemeral_port, DEFAULT_BASE_PORT + 1);
assert_eq!(w1.cpu_weight, 100);
assert_eq!(w1.memory_limit_bytes, 2 * 1024 * 1024 * 1024);
assert_eq!(w1.pids_max, 128);
assert!(w1.ipc_isolated);
assert!(fleet.verify_isolation("agent-01", "agent-02").is_ok());
assert_eq!(fleet.active_count(), 2);
fleet.release_worker("agent-01").expect("release w1");
assert_eq!(fleet.active_count(), 1);
let w3 = fleet.allocate_worker("codex-1").expect("allocate w3");
assert_eq!(w3.worker_id, "agent-01");
}
#[test]
fn test_fleet_capacity_limit() {
let config = FleetConfig {
max_agents: 3,
..Default::default()
};
let fleet = FleetManager::new(config);
let _w1 = fleet.allocate_worker("a1").unwrap();
let _w2 = fleet.allocate_worker("a2").unwrap();
let _w3 = fleet.allocate_worker("a3").unwrap();
let overflow = fleet.allocate_worker("a4");
assert!(overflow.is_err());
assert!(overflow
.unwrap_err()
.to_string()
.contains("Fleet capacity exceeded"));
}
}