use crate::tools::{policy_sandbox::PolicyConfig, ToolRegistry};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum SandboxMode {
#[default]
None,
Policy,
Container,
MicroVm,
}
pub trait ToolSetProvider: Send + Sync + 'static {
fn build_registry(&self) -> ToolRegistry;
fn sandbox_mode(&self) -> SandboxMode;
}
pub struct LocalToolSetProvider {
workspace: std::path::PathBuf,
shell_timeout_secs: u64,
skills: Vec<crate::skills::Skill>,
}
impl LocalToolSetProvider {
pub fn new(
workspace: std::path::PathBuf,
shell_timeout_secs: u64,
skills: Vec<crate::skills::Skill>,
) -> Self {
Self {
workspace,
shell_timeout_secs,
skills,
}
}
}
impl ToolSetProvider for LocalToolSetProvider {
fn build_registry(&self) -> ToolRegistry {
crate::tools::build_standard_tools(&self.workspace, &self.skills, self.shell_timeout_secs)
}
fn sandbox_mode(&self) -> SandboxMode {
SandboxMode::None
}
}
pub struct PolicyToolSetProvider {
workspace: std::path::PathBuf,
shell_timeout_secs: u64,
skills: Vec<crate::skills::Skill>,
pub policy: PolicyConfig,
}
impl PolicyToolSetProvider {
pub fn new(
workspace: std::path::PathBuf,
shell_timeout_secs: u64,
skills: Vec<crate::skills::Skill>,
policy: PolicyConfig,
) -> Self {
Self {
workspace,
shell_timeout_secs,
skills,
policy,
}
}
pub fn restrictive(
workspace: std::path::PathBuf,
shell_timeout_secs: u64,
skills: Vec<crate::skills::Skill>,
) -> Self {
Self::new(
workspace,
shell_timeout_secs,
skills,
PolicyConfig::default_restrictive(),
)
}
}
impl ToolSetProvider for PolicyToolSetProvider {
fn build_registry(&self) -> ToolRegistry {
crate::tools::build_standard_tools(&self.workspace, &self.skills, self.shell_timeout_secs)
.with_policy(self.policy.clone())
}
fn sandbox_mode(&self) -> SandboxMode {
SandboxMode::Policy
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
fn provider() -> (LocalToolSetProvider, TempDir) {
let dir = TempDir::new().unwrap();
let p = LocalToolSetProvider::new(dir.path().to_path_buf(), 30, vec![]);
(p, dir)
}
#[test]
fn local_provider_sandbox_mode_is_none() {
let (p, _dir) = provider();
assert_eq!(p.sandbox_mode(), SandboxMode::None);
}
#[test]
fn local_provider_registry_non_empty() {
let (p, _dir) = provider();
let reg = p.build_registry();
let names = reg.names();
assert!(
names.iter().any(|n| n == "run_shell"),
"expected run_shell in registry, got: {names:?}"
);
}
#[test]
fn policy_provider_sandbox_mode_is_policy() {
let dir = TempDir::new().unwrap();
let p = PolicyToolSetProvider::restrictive(dir.path().to_path_buf(), 30, vec![]);
assert_eq!(p.sandbox_mode(), SandboxMode::Policy);
}
#[test]
fn policy_provider_attaches_policy_to_registry() {
let dir = TempDir::new().unwrap();
let p = PolicyToolSetProvider::restrictive(dir.path().to_path_buf(), 30, vec![]);
let reg = p.build_registry();
let policy = reg.policy().expect("policy should be attached");
assert!(policy.check_shell("ls").is_ok());
assert!(policy.check_shell("rm -rf /").is_err());
}
}