mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
Merge remote-tracking branch 'origin/main'
This commit is contained in:
+15
-14
@@ -64,13 +64,18 @@ pub struct Agent {
|
||||
|
||||
impl Agent {
|
||||
/// Create a new agent.
|
||||
///
|
||||
/// Optionally accepts a pre-created `ContextManager` for sharing with job tools.
|
||||
/// If not provided, creates a new one.
|
||||
pub fn new(
|
||||
config: AgentConfig,
|
||||
deps: AgentDeps,
|
||||
channels: ChannelManager,
|
||||
heartbeat_config: Option<HeartbeatConfig>,
|
||||
context_manager: Option<Arc<ContextManager>>,
|
||||
) -> Self {
|
||||
let context_manager = Arc::new(ContextManager::new(config.max_parallel_jobs));
|
||||
let context_manager = context_manager
|
||||
.unwrap_or_else(|| Arc::new(ContextManager::new(config.max_parallel_jobs)));
|
||||
|
||||
let scheduler = Arc::new(Scheduler::new(
|
||||
config.clone(),
|
||||
@@ -388,25 +393,21 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
|
||||
// Route for job commands (bypass turn system)
|
||||
// Build a temporary message with the content to route
|
||||
// Handle explicit commands (starting with /) directly
|
||||
// Everything else goes through the normal agentic loop with tools
|
||||
let temp_message = IncomingMessage {
|
||||
content: content.to_string(),
|
||||
..message.clone()
|
||||
};
|
||||
let intent = self.router.route(&temp_message);
|
||||
match &intent {
|
||||
MessageIntent::CreateJob { .. }
|
||||
| MessageIntent::CheckJobStatus { .. }
|
||||
| MessageIntent::CancelJob { .. }
|
||||
| MessageIntent::ListJobs { .. }
|
||||
| MessageIntent::HelpJob { .. }
|
||||
| MessageIntent::Command { .. } => {
|
||||
return self.handle_job_or_command(intent, message).await;
|
||||
}
|
||||
_ => {}
|
||||
|
||||
if let Some(intent) = self.router.route_command(&temp_message) {
|
||||
// Explicit command like /status, /job, /list - handle directly
|
||||
return self.handle_job_or_command(intent, message).await;
|
||||
}
|
||||
|
||||
// Natural language goes through the agentic loop
|
||||
// Job tools (create_job, list_jobs, etc.) are in the tool registry
|
||||
|
||||
// Auto-compact if needed BEFORE adding new turn
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
|
||||
+59
-123
@@ -1,4 +1,8 @@
|
||||
//! Message routing to appropriate handlers.
|
||||
//!
|
||||
//! The router handles explicit commands (starting with `/`).
|
||||
//! Natural language intent classification is handled by `IntentClassifier`
|
||||
//! which uses LLM + tools instead of brittle pattern matching.
|
||||
|
||||
use crate::channels::IncomingMessage;
|
||||
|
||||
@@ -27,7 +31,9 @@ pub enum MessageIntent {
|
||||
Unknown,
|
||||
}
|
||||
|
||||
/// Routes messages to appropriate handlers based on intent.
|
||||
/// Routes messages to appropriate handlers based on explicit commands.
|
||||
///
|
||||
/// For natural language messages, use `IntentClassifier` instead.
|
||||
pub struct Router {
|
||||
/// Command prefix (e.g., "/" or "!")
|
||||
command_prefix: String,
|
||||
@@ -47,17 +53,23 @@ impl Router {
|
||||
self
|
||||
}
|
||||
|
||||
/// Route a message to determine its intent.
|
||||
pub fn route(&self, message: &IncomingMessage) -> MessageIntent {
|
||||
/// Check if a message is an explicit command.
|
||||
pub fn is_command(&self, message: &IncomingMessage) -> bool {
|
||||
message.content.trim().starts_with(&self.command_prefix)
|
||||
}
|
||||
|
||||
/// Route an explicit command to determine its intent.
|
||||
///
|
||||
/// Returns `None` if the message is not a command.
|
||||
/// For non-commands, use `IntentClassifier::classify()` instead.
|
||||
pub fn route_command(&self, message: &IncomingMessage) -> Option<MessageIntent> {
|
||||
let content = message.content.trim();
|
||||
|
||||
// Check for commands
|
||||
if content.starts_with(&self.command_prefix) {
|
||||
return self.parse_command(content);
|
||||
Some(self.parse_command(content))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
|
||||
// Try to extract intent from natural language
|
||||
self.extract_intent(content)
|
||||
}
|
||||
|
||||
fn parse_command(&self, content: &str) -> MessageIntent {
|
||||
@@ -111,61 +123,6 @@ impl Router {
|
||||
None => MessageIntent::Unknown,
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_intent(&self, content: &str) -> MessageIntent {
|
||||
let lower = content.to_lowercase();
|
||||
|
||||
// Job creation patterns - must be explicit about creating a job
|
||||
// More specific patterns to avoid capturing general conversation
|
||||
let is_job_creation = lower.starts_with("create job ")
|
||||
|| lower.starts_with("new job ")
|
||||
|| lower.starts_with("schedule job ")
|
||||
|| lower.starts_with("run job ")
|
||||
|| (lower.contains("create") && lower.contains("job"));
|
||||
|
||||
if is_job_creation {
|
||||
return MessageIntent::CreateJob {
|
||||
title: extract_title(content),
|
||||
description: content.to_string(),
|
||||
category: extract_category(content),
|
||||
};
|
||||
}
|
||||
|
||||
// Status check patterns
|
||||
if lower.contains("status")
|
||||
|| lower.contains("how is")
|
||||
|| lower.contains("progress")
|
||||
|| lower.starts_with("check ")
|
||||
{
|
||||
return MessageIntent::CheckJobStatus {
|
||||
job_id: extract_job_id(content),
|
||||
};
|
||||
}
|
||||
|
||||
// Cancel patterns
|
||||
if lower.contains("cancel") || lower.contains("stop") || lower.contains("abort") {
|
||||
if let Some(job_id) = extract_job_id(content) {
|
||||
return MessageIntent::CancelJob { job_id };
|
||||
}
|
||||
}
|
||||
|
||||
// List patterns
|
||||
if lower.starts_with("list") || lower.contains("show jobs") || lower.contains("my jobs") {
|
||||
return MessageIntent::ListJobs { filter: None };
|
||||
}
|
||||
|
||||
// Help patterns
|
||||
if lower.contains("stuck") || lower.contains("not working") || lower.contains("fix") {
|
||||
if let Some(job_id) = extract_job_id(content) {
|
||||
return MessageIntent::HelpJob { job_id };
|
||||
}
|
||||
}
|
||||
|
||||
// Default to chat
|
||||
MessageIntent::Chat {
|
||||
content: content.to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for Router {
|
||||
@@ -174,53 +131,6 @@ impl Default for Router {
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract a title from content.
|
||||
fn extract_title(content: &str) -> String {
|
||||
// Take first sentence or first N characters
|
||||
let first_sentence = content.split('.').next().unwrap_or(content);
|
||||
let title = first_sentence.chars().take(100).collect::<String>();
|
||||
if title.len() < first_sentence.len() {
|
||||
format!("{}...", title)
|
||||
} else {
|
||||
title
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract a category from content.
|
||||
fn extract_category(content: &str) -> Option<String> {
|
||||
let lower = content.to_lowercase();
|
||||
|
||||
let categories = [
|
||||
("code", "development"),
|
||||
("program", "development"),
|
||||
("website", "web"),
|
||||
("api", "development"),
|
||||
("data", "data"),
|
||||
("write", "writing"),
|
||||
("design", "design"),
|
||||
("research", "research"),
|
||||
];
|
||||
|
||||
for (keyword, category) in categories {
|
||||
if lower.contains(keyword) {
|
||||
return Some(category.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
/// Extract a job ID from content.
|
||||
fn extract_job_id(content: &str) -> Option<String> {
|
||||
// Look for UUID patterns
|
||||
let uuid_regex = regex::Regex::new(
|
||||
r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}",
|
||||
)
|
||||
.ok()?;
|
||||
|
||||
uuid_regex.find(content).map(|m| m.as_str().to_string())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -230,35 +140,61 @@ mod tests {
|
||||
let router = Router::new();
|
||||
|
||||
let msg = IncomingMessage::new("test", "user", "/status abc-123");
|
||||
let intent = router.route(&msg);
|
||||
let intent = router.route_command(&msg);
|
||||
|
||||
assert!(matches!(intent, MessageIntent::CheckJobStatus { .. }));
|
||||
assert!(matches!(intent, Some(MessageIntent::CheckJobStatus { .. })));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_natural_language_routing() {
|
||||
fn test_is_command() {
|
||||
let router = Router::new();
|
||||
|
||||
let cmd_msg = IncomingMessage::new("test", "user", "/status");
|
||||
assert!(router.is_command(&cmd_msg));
|
||||
|
||||
let chat_msg = IncomingMessage::new("test", "user", "Hello there");
|
||||
assert!(!router.is_command(&chat_msg));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_non_command_returns_none() {
|
||||
let router = Router::new();
|
||||
|
||||
// Natural language messages return None - they should use IntentClassifier
|
||||
let msg = IncomingMessage::new("test", "user", "Can you create a website for me?");
|
||||
let intent = router.route(&msg);
|
||||
assert!(router.route_command(&msg).is_none());
|
||||
|
||||
assert!(matches!(intent, MessageIntent::CreateJob { .. }));
|
||||
let msg2 = IncomingMessage::new("test", "user", "Hello, how are you?");
|
||||
assert!(router.route_command(&msg2).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_chat_fallback() {
|
||||
fn test_command_create_job() {
|
||||
let router = Router::new();
|
||||
|
||||
let msg = IncomingMessage::new("test", "user", "Hello, how are you?");
|
||||
let intent = router.route(&msg);
|
||||
let msg = IncomingMessage::new("test", "user", "/job build a website");
|
||||
let intent = router.route_command(&msg);
|
||||
|
||||
assert!(matches!(intent, MessageIntent::Chat { .. }));
|
||||
match intent {
|
||||
Some(MessageIntent::CreateJob { title, .. }) => {
|
||||
assert_eq!(title, "build a website");
|
||||
}
|
||||
_ => panic!("Expected CreateJob intent"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_job_id() {
|
||||
let content = "Check status of job 550e8400-e29b-41d4-a716-446655440000";
|
||||
let id = extract_job_id(content);
|
||||
assert_eq!(id, Some("550e8400-e29b-41d4-a716-446655440000".to_string()));
|
||||
fn test_command_list_jobs() {
|
||||
let router = Router::new();
|
||||
|
||||
let msg = IncomingMessage::new("test", "user", "/list active");
|
||||
let intent = router.route_command(&msg);
|
||||
|
||||
match intent {
|
||||
Some(MessageIntent::ListJobs { filter }) => {
|
||||
assert_eq!(filter, Some("active".to_string()));
|
||||
}
|
||||
_ => panic!("Expected ListJobs intent"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+97
-2
@@ -21,6 +21,7 @@ pub struct Config {
|
||||
pub secrets: SecretsConfig,
|
||||
pub builder: BuilderModeConfig,
|
||||
pub heartbeat: HeartbeatConfig,
|
||||
pub sandbox: SandboxModeConfig,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
@@ -41,6 +42,7 @@ impl Config {
|
||||
secrets: SecretsConfig::from_env()?,
|
||||
builder: BuilderModeConfig::from_env()?,
|
||||
heartbeat: HeartbeatConfig::from_env()?,
|
||||
sandbox: SandboxModeConfig::from_env()?,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -599,7 +601,7 @@ pub struct BuilderModeConfig {
|
||||
impl Default for BuilderModeConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: false,
|
||||
enabled: true, // Builder enabled by default
|
||||
build_dir: None,
|
||||
max_iterations: 20,
|
||||
timeout_secs: 600,
|
||||
@@ -618,7 +620,7 @@ impl BuilderModeConfig {
|
||||
key: "BUILDER_ENABLED".to_string(),
|
||||
message: format!("must be 'true' or 'false': {e}"),
|
||||
})?
|
||||
.unwrap_or(false),
|
||||
.unwrap_or(true), // Builder enabled by default
|
||||
build_dir: optional_env("BUILDER_DIR")?.map(PathBuf::from),
|
||||
max_iterations: parse_optional_env("BUILDER_MAX_ITERATIONS", 20)?,
|
||||
timeout_secs: parse_optional_env("BUILDER_TIMEOUT_SECS", 600)?,
|
||||
@@ -690,6 +692,99 @@ impl HeartbeatConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Docker sandbox configuration.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SandboxModeConfig {
|
||||
/// Whether the Docker sandbox is enabled.
|
||||
pub enabled: bool,
|
||||
/// Sandbox policy: "readonly", "workspace_write", or "full_access".
|
||||
pub policy: String,
|
||||
/// Command timeout in seconds.
|
||||
pub timeout_secs: u64,
|
||||
/// Memory limit in megabytes.
|
||||
pub memory_limit_mb: u64,
|
||||
/// CPU shares (relative weight).
|
||||
pub cpu_shares: u32,
|
||||
/// Docker image for the sandbox.
|
||||
pub image: String,
|
||||
/// Whether to auto-pull the image if not found.
|
||||
pub auto_pull_image: bool,
|
||||
/// Additional domains to allow through the network proxy.
|
||||
pub extra_allowed_domains: Vec<String>,
|
||||
}
|
||||
|
||||
impl Default for SandboxModeConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: false, // Disabled by default
|
||||
policy: "readonly".to_string(),
|
||||
timeout_secs: 120,
|
||||
memory_limit_mb: 2048,
|
||||
cpu_shares: 1024,
|
||||
image: "ghcr.io/nearai/sandbox:latest".to_string(),
|
||||
auto_pull_image: true,
|
||||
extra_allowed_domains: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SandboxModeConfig {
|
||||
fn from_env() -> Result<Self, ConfigError> {
|
||||
let extra_domains = optional_env("SANDBOX_EXTRA_DOMAINS")?
|
||||
.map(|s| s.split(',').map(|d| d.trim().to_string()).collect())
|
||||
.unwrap_or_default();
|
||||
|
||||
Ok(Self {
|
||||
enabled: optional_env("SANDBOX_ENABLED")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
.map_err(|e| ConfigError::InvalidValue {
|
||||
key: "SANDBOX_ENABLED".to_string(),
|
||||
message: format!("must be 'true' or 'false': {e}"),
|
||||
})?
|
||||
.unwrap_or(true),
|
||||
policy: optional_env("SANDBOX_POLICY")?.unwrap_or_else(|| "readonly".to_string()),
|
||||
timeout_secs: parse_optional_env("SANDBOX_TIMEOUT_SECS", 120)?,
|
||||
memory_limit_mb: parse_optional_env("SANDBOX_MEMORY_LIMIT_MB", 2048)?,
|
||||
cpu_shares: parse_optional_env("SANDBOX_CPU_SHARES", 1024)?,
|
||||
image: optional_env("SANDBOX_IMAGE")?
|
||||
.unwrap_or_else(|| "ghcr.io/nearai/sandbox:latest".to_string()),
|
||||
auto_pull_image: optional_env("SANDBOX_AUTO_PULL")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
.map_err(|e| ConfigError::InvalidValue {
|
||||
key: "SANDBOX_AUTO_PULL".to_string(),
|
||||
message: format!("must be 'true' or 'false': {e}"),
|
||||
})?
|
||||
.unwrap_or(true),
|
||||
extra_allowed_domains: extra_domains,
|
||||
})
|
||||
}
|
||||
|
||||
/// Convert to SandboxConfig for the sandbox module.
|
||||
pub fn to_sandbox_config(&self) -> crate::sandbox::SandboxConfig {
|
||||
use crate::sandbox::SandboxPolicy;
|
||||
use std::time::Duration;
|
||||
|
||||
let policy = self.policy.parse().unwrap_or(SandboxPolicy::ReadOnly);
|
||||
|
||||
let mut allowlist = crate::sandbox::default_allowlist();
|
||||
allowlist.extend(self.extra_allowed_domains.clone());
|
||||
|
||||
crate::sandbox::SandboxConfig {
|
||||
enabled: self.enabled,
|
||||
policy,
|
||||
timeout: Duration::from_secs(self.timeout_secs),
|
||||
memory_limit_mb: self.memory_limit_mb,
|
||||
cpu_shares: self.cpu_shares,
|
||||
network_allowlist: allowlist,
|
||||
image: self.image.clone(),
|
||||
auto_pull_image: self.auto_pull_image,
|
||||
proxy_port: 0, // Auto-assign
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Helper functions
|
||||
|
||||
fn required_env(key: &str) -> Result<String, ConfigError> {
|
||||
|
||||
@@ -49,6 +49,7 @@ pub mod evaluation;
|
||||
pub mod history;
|
||||
pub mod llm;
|
||||
pub mod safety;
|
||||
pub mod sandbox;
|
||||
pub mod secrets;
|
||||
pub mod settings;
|
||||
pub mod setup;
|
||||
|
||||
@@ -16,6 +16,7 @@ use near_agent::{
|
||||
},
|
||||
cli::{Cli, Command, run_tool_command},
|
||||
config::Config,
|
||||
context::ContextManager,
|
||||
history::Store,
|
||||
llm::{SessionConfig, create_llm_provider, create_session_manager},
|
||||
safety::SafetyLayer,
|
||||
@@ -509,6 +510,12 @@ async fn main() -> anyhow::Result<()> {
|
||||
}
|
||||
}
|
||||
|
||||
// Create context manager (shared between job tools and agent)
|
||||
let context_manager = Arc::new(ContextManager::new(config.agent.max_parallel_jobs));
|
||||
|
||||
// Register job tools
|
||||
tools.register_job_tools(Arc::clone(&context_manager));
|
||||
|
||||
// Create and run the agent
|
||||
let deps = AgentDeps {
|
||||
store,
|
||||
@@ -522,6 +529,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
deps,
|
||||
channels,
|
||||
Some(config.heartbeat.clone()),
|
||||
Some(context_manager),
|
||||
);
|
||||
|
||||
tracing::info!("Agent initialized, starting main loop...");
|
||||
|
||||
@@ -0,0 +1,258 @@
|
||||
//! Configuration for the Docker execution sandbox.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
/// Configuration for the sandbox system.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SandboxConfig {
|
||||
/// Whether the sandbox is enabled.
|
||||
pub enabled: bool,
|
||||
/// Security policy for sandbox execution.
|
||||
pub policy: SandboxPolicy,
|
||||
/// Default timeout for command execution.
|
||||
pub timeout: Duration,
|
||||
/// Memory limit in megabytes.
|
||||
pub memory_limit_mb: u64,
|
||||
/// CPU shares (relative weight, default 1024).
|
||||
pub cpu_shares: u32,
|
||||
/// Network allowlist for proxied requests.
|
||||
pub network_allowlist: Vec<String>,
|
||||
/// Docker image to use for the sandbox.
|
||||
pub image: String,
|
||||
/// Whether to auto-pull the image if not found.
|
||||
pub auto_pull_image: bool,
|
||||
/// Port for the HTTP proxy (0 = auto-assign).
|
||||
pub proxy_port: u16,
|
||||
}
|
||||
|
||||
impl Default for SandboxConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: false, // Disabled by default until Docker is confirmed available
|
||||
policy: SandboxPolicy::ReadOnly,
|
||||
timeout: Duration::from_secs(120),
|
||||
memory_limit_mb: 2048,
|
||||
cpu_shares: 1024,
|
||||
network_allowlist: default_allowlist(),
|
||||
image: "ghcr.io/nearai/sandbox:latest".to_string(),
|
||||
auto_pull_image: true,
|
||||
proxy_port: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Security policy for sandbox execution.
|
||||
///
|
||||
/// ```text
|
||||
/// ┌─────────────────────────────────────────────────────────────────────┐
|
||||
/// │ Sandbox Policies │
|
||||
/// ├─────────────────┬──────────────────┬────────────────────────────────┤
|
||||
/// │ Policy │ Filesystem │ Network │
|
||||
/// ├─────────────────┼──────────────────┼────────────────────────────────┤
|
||||
/// │ ReadOnly │ /workspace (ro) │ Proxied (allowlist only) │
|
||||
/// │ WorkspaceWrite │ /workspace (rw) │ Proxied (allowlist only) │
|
||||
/// │ FullAccess │ Full host │ Full network (DANGER) │
|
||||
/// └─────────────────┴──────────────────┴────────────────────────────────┘
|
||||
/// ```
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum SandboxPolicy {
|
||||
/// Read-only access to workspace, proxied network.
|
||||
/// Use for: exploring code, fetching docs, read-only operations.
|
||||
#[default]
|
||||
ReadOnly,
|
||||
|
||||
/// Read/write access to workspace, proxied network.
|
||||
/// Use for: building software, running tests, generating files.
|
||||
WorkspaceWrite,
|
||||
|
||||
/// Full access (no sandbox). Use with extreme caution.
|
||||
/// This bypasses all isolation and runs directly on host.
|
||||
FullAccess,
|
||||
}
|
||||
|
||||
impl SandboxPolicy {
|
||||
/// Returns true if filesystem writes are allowed.
|
||||
pub fn allows_writes(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
SandboxPolicy::WorkspaceWrite | SandboxPolicy::FullAccess
|
||||
)
|
||||
}
|
||||
|
||||
/// Returns true if network requests bypass the proxy.
|
||||
pub fn has_full_network(&self) -> bool {
|
||||
matches!(self, SandboxPolicy::FullAccess)
|
||||
}
|
||||
|
||||
/// Returns true if running in a container.
|
||||
pub fn is_sandboxed(&self) -> bool {
|
||||
!matches!(self, SandboxPolicy::FullAccess)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::str::FromStr for SandboxPolicy {
|
||||
type Err = String;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
match s.to_lowercase().as_str() {
|
||||
"readonly" | "read_only" | "ro" => Ok(SandboxPolicy::ReadOnly),
|
||||
"workspacewrite" | "workspace_write" | "rw" => Ok(SandboxPolicy::WorkspaceWrite),
|
||||
"fullaccess" | "full_access" | "full" | "none" => Ok(SandboxPolicy::FullAccess),
|
||||
_ => Err(format!(
|
||||
"invalid sandbox policy '{}', expected 'readonly', 'workspace_write', or 'full_access'",
|
||||
s
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Resource limits for container execution.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ResourceLimits {
|
||||
/// Maximum memory in bytes.
|
||||
pub memory_bytes: u64,
|
||||
/// CPU shares (relative weight).
|
||||
pub cpu_shares: u32,
|
||||
/// Maximum execution time.
|
||||
pub timeout: Duration,
|
||||
/// Maximum output size in bytes.
|
||||
pub max_output_bytes: usize,
|
||||
}
|
||||
|
||||
impl Default for ResourceLimits {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
memory_bytes: 2 * 1024 * 1024 * 1024, // 2 GB
|
||||
cpu_shares: 1024,
|
||||
timeout: Duration::from_secs(120),
|
||||
max_output_bytes: 64 * 1024, // 64 KB
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Default network allowlist for common development operations.
|
||||
pub fn default_allowlist() -> Vec<String> {
|
||||
vec![
|
||||
// Package registries
|
||||
"crates.io".to_string(),
|
||||
"static.crates.io".to_string(),
|
||||
"index.crates.io".to_string(),
|
||||
"registry.npmjs.org".to_string(),
|
||||
"proxy.golang.org".to_string(),
|
||||
"pypi.org".to_string(),
|
||||
"files.pythonhosted.org".to_string(),
|
||||
// Documentation
|
||||
"docs.rs".to_string(),
|
||||
"doc.rust-lang.org".to_string(),
|
||||
"nodejs.org".to_string(),
|
||||
"go.dev".to_string(),
|
||||
"docs.python.org".to_string(),
|
||||
// Version control (read-only)
|
||||
"github.com".to_string(),
|
||||
"raw.githubusercontent.com".to_string(),
|
||||
"api.github.com".to_string(),
|
||||
"codeload.github.com".to_string(),
|
||||
// Common APIs (credentials will be injected by proxy)
|
||||
"api.openai.com".to_string(),
|
||||
"api.anthropic.com".to_string(),
|
||||
"api.near.ai".to_string(),
|
||||
]
|
||||
}
|
||||
|
||||
/// Credential injection configuration.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CredentialMapping {
|
||||
/// Domain this credential applies to.
|
||||
pub domain: String,
|
||||
/// Name of the secret to inject.
|
||||
pub secret_name: String,
|
||||
/// Where to inject the credential.
|
||||
pub location: CredentialLocation,
|
||||
}
|
||||
|
||||
/// Where to inject a credential in an HTTP request.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum CredentialLocation {
|
||||
/// Inject as Authorization: Bearer <token>
|
||||
AuthorizationBearer,
|
||||
/// Inject as a custom header.
|
||||
Header(String),
|
||||
/// Inject as a query parameter.
|
||||
QueryParam(String),
|
||||
}
|
||||
|
||||
impl Default for CredentialMapping {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
domain: String::new(),
|
||||
secret_name: String::new(),
|
||||
location: CredentialLocation::AuthorizationBearer,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Default credential mappings for common APIs.
|
||||
pub fn default_credential_mappings() -> Vec<CredentialMapping> {
|
||||
vec![
|
||||
CredentialMapping {
|
||||
domain: "api.openai.com".to_string(),
|
||||
secret_name: "OPENAI_API_KEY".to_string(),
|
||||
location: CredentialLocation::AuthorizationBearer,
|
||||
},
|
||||
CredentialMapping {
|
||||
domain: "api.anthropic.com".to_string(),
|
||||
secret_name: "ANTHROPIC_API_KEY".to_string(),
|
||||
location: CredentialLocation::Header("x-api-key".to_string()),
|
||||
},
|
||||
CredentialMapping {
|
||||
domain: "api.near.ai".to_string(),
|
||||
secret_name: "NEARAI_API_KEY".to_string(),
|
||||
location: CredentialLocation::AuthorizationBearer,
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_policy_parsing() {
|
||||
assert_eq!(
|
||||
"readonly".parse::<SandboxPolicy>().unwrap(),
|
||||
SandboxPolicy::ReadOnly
|
||||
);
|
||||
assert_eq!(
|
||||
"workspace_write".parse::<SandboxPolicy>().unwrap(),
|
||||
SandboxPolicy::WorkspaceWrite
|
||||
);
|
||||
assert_eq!(
|
||||
"full_access".parse::<SandboxPolicy>().unwrap(),
|
||||
SandboxPolicy::FullAccess
|
||||
);
|
||||
assert!("invalid".parse::<SandboxPolicy>().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_policy_properties() {
|
||||
assert!(!SandboxPolicy::ReadOnly.allows_writes());
|
||||
assert!(SandboxPolicy::WorkspaceWrite.allows_writes());
|
||||
assert!(SandboxPolicy::FullAccess.allows_writes());
|
||||
|
||||
assert!(!SandboxPolicy::ReadOnly.has_full_network());
|
||||
assert!(!SandboxPolicy::WorkspaceWrite.has_full_network());
|
||||
assert!(SandboxPolicy::FullAccess.has_full_network());
|
||||
|
||||
assert!(SandboxPolicy::ReadOnly.is_sandboxed());
|
||||
assert!(SandboxPolicy::WorkspaceWrite.is_sandboxed());
|
||||
assert!(!SandboxPolicy::FullAccess.is_sandboxed());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_allowlist_has_common_registries() {
|
||||
let allowlist = default_allowlist();
|
||||
assert!(allowlist.contains(&"crates.io".to_string()));
|
||||
assert!(allowlist.contains(&"registry.npmjs.org".to_string()));
|
||||
assert!(allowlist.contains(&"github.com".to_string()));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,519 @@
|
||||
//! Docker container lifecycle management.
|
||||
//!
|
||||
//! Handles creating, running, and cleaning up containers for sandboxed execution.
|
||||
//!
|
||||
//! # Container Setup
|
||||
//!
|
||||
//! ```text
|
||||
//! ┌────────────────────────────────────────────────────────────────────────┐
|
||||
//! │ Docker Container │
|
||||
//! │ │
|
||||
//! │ Environment: │
|
||||
//! │ http_proxy=http://host.docker.internal:PORT │
|
||||
//! │ https_proxy=http://host.docker.internal:PORT │
|
||||
//! │ (No secrets or credentials) │
|
||||
//! │ │
|
||||
//! │ Mounts: │
|
||||
//! │ /workspace ─▶ Host working directory (ro or rw based on policy) │
|
||||
//! │ /output ─▶ Output directory for artifacts (rw) │
|
||||
//! │ │
|
||||
//! │ Limits: │
|
||||
//! │ Memory: 2GB (default) │
|
||||
//! │ CPU: 1024 shares │
|
||||
//! │ No privileged mode │
|
||||
//! │ Non-root user (UID 1000) │
|
||||
//! └────────────────────────────────────────────────────────────────────────┘
|
||||
//! ```
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
|
||||
use bollard::Docker;
|
||||
use bollard::container::{
|
||||
Config, CreateContainerOptions, LogOutput, LogsOptions, RemoveContainerOptions,
|
||||
StartContainerOptions, WaitContainerOptions,
|
||||
};
|
||||
use bollard::exec::{CreateExecOptions, StartExecResults};
|
||||
use bollard::models::HostConfig;
|
||||
use futures::StreamExt;
|
||||
|
||||
use crate::sandbox::config::{ResourceLimits, SandboxPolicy};
|
||||
use crate::sandbox::error::{Result, SandboxError};
|
||||
|
||||
/// Output from container execution.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ContainerOutput {
|
||||
/// Exit code from the command.
|
||||
pub exit_code: i64,
|
||||
/// Standard output.
|
||||
pub stdout: String,
|
||||
/// Standard error.
|
||||
pub stderr: String,
|
||||
/// How long the command ran.
|
||||
pub duration: Duration,
|
||||
/// Whether output was truncated.
|
||||
pub truncated: bool,
|
||||
}
|
||||
|
||||
/// Manages Docker container lifecycle.
|
||||
pub struct ContainerRunner {
|
||||
docker: Docker,
|
||||
image: String,
|
||||
proxy_port: u16,
|
||||
}
|
||||
|
||||
impl ContainerRunner {
|
||||
/// Create a new container runner.
|
||||
pub fn new(docker: Docker, image: String, proxy_port: u16) -> Self {
|
||||
Self {
|
||||
docker,
|
||||
image,
|
||||
proxy_port,
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if the Docker daemon is available.
|
||||
pub async fn is_available(&self) -> bool {
|
||||
self.docker.ping().await.is_ok()
|
||||
}
|
||||
|
||||
/// Check if the sandbox image exists locally.
|
||||
pub async fn image_exists(&self) -> bool {
|
||||
self.docker.inspect_image(&self.image).await.is_ok()
|
||||
}
|
||||
|
||||
/// Pull the sandbox image.
|
||||
pub async fn pull_image(&self) -> Result<()> {
|
||||
use bollard::image::CreateImageOptions;
|
||||
|
||||
tracing::info!("Pulling sandbox image: {}", self.image);
|
||||
|
||||
let options = CreateImageOptions {
|
||||
from_image: self.image.clone(),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let mut stream = self.docker.create_image(Some(options), None, None);
|
||||
|
||||
while let Some(result) = stream.next().await {
|
||||
match result {
|
||||
Ok(info) => {
|
||||
if let Some(status) = info.status {
|
||||
tracing::debug!("Pull status: {}", status);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(SandboxError::ContainerCreationFailed {
|
||||
reason: format!("image pull failed: {}", e),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!("Successfully pulled image: {}", self.image);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Execute a command in a new container.
|
||||
pub async fn execute(
|
||||
&self,
|
||||
command: &str,
|
||||
working_dir: &Path,
|
||||
policy: SandboxPolicy,
|
||||
limits: &ResourceLimits,
|
||||
env: HashMap<String, String>,
|
||||
) -> Result<ContainerOutput> {
|
||||
let start_time = std::time::Instant::now();
|
||||
|
||||
// Create the container
|
||||
let container_id = self
|
||||
.create_container(command, working_dir, policy, limits, env)
|
||||
.await?;
|
||||
|
||||
// Start the container
|
||||
self.docker
|
||||
.start_container(&container_id, None::<StartContainerOptions<String>>)
|
||||
.await
|
||||
.map_err(|e| SandboxError::ContainerStartFailed {
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
// Wait for completion with timeout
|
||||
let result = tokio::time::timeout(limits.timeout, async {
|
||||
self.wait_for_container(&container_id, limits.max_output_bytes)
|
||||
.await
|
||||
})
|
||||
.await;
|
||||
|
||||
// Always clean up the container
|
||||
let _ = self
|
||||
.docker
|
||||
.remove_container(
|
||||
&container_id,
|
||||
Some(RemoveContainerOptions {
|
||||
force: true,
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(mut output)) => {
|
||||
output.duration = start_time.elapsed();
|
||||
Ok(output)
|
||||
}
|
||||
Ok(Err(e)) => Err(e),
|
||||
Err(_) => Err(SandboxError::Timeout(limits.timeout)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Execute a command in an existing container using exec.
|
||||
pub async fn exec_in_container(
|
||||
&self,
|
||||
container_id: &str,
|
||||
command: &str,
|
||||
working_dir: &str,
|
||||
limits: &ResourceLimits,
|
||||
) -> Result<ContainerOutput> {
|
||||
let start_time = std::time::Instant::now();
|
||||
|
||||
let exec = self
|
||||
.docker
|
||||
.create_exec(
|
||||
container_id,
|
||||
CreateExecOptions {
|
||||
cmd: Some(vec!["sh", "-c", command]),
|
||||
attach_stdout: Some(true),
|
||||
attach_stderr: Some(true),
|
||||
working_dir: Some(working_dir),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map_err(|e| SandboxError::ExecutionFailed {
|
||||
reason: format!("exec create failed: {}", e),
|
||||
})?;
|
||||
|
||||
let result = tokio::time::timeout(
|
||||
limits.timeout,
|
||||
self.run_exec(&exec.id, limits.max_output_bytes),
|
||||
)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(mut output)) => {
|
||||
output.duration = start_time.elapsed();
|
||||
Ok(output)
|
||||
}
|
||||
Ok(Err(e)) => Err(e),
|
||||
Err(_) => Err(SandboxError::Timeout(limits.timeout)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a container with the appropriate configuration.
|
||||
async fn create_container(
|
||||
&self,
|
||||
command: &str,
|
||||
working_dir: &Path,
|
||||
policy: SandboxPolicy,
|
||||
limits: &ResourceLimits,
|
||||
env: HashMap<String, String>,
|
||||
) -> Result<String> {
|
||||
let working_dir_str = working_dir.display().to_string();
|
||||
|
||||
// Build environment variables
|
||||
let mut env_vec: Vec<String> = env
|
||||
.into_iter()
|
||||
.map(|(k, v)| format!("{}={}", k, v))
|
||||
.collect();
|
||||
|
||||
// Add proxy environment (uses host.docker.internal for Mac/Windows, 172.17.0.1 for Linux)
|
||||
let proxy_host = if cfg!(target_os = "linux") {
|
||||
"172.17.0.1"
|
||||
} else {
|
||||
"host.docker.internal"
|
||||
};
|
||||
|
||||
if self.proxy_port > 0 && policy.is_sandboxed() {
|
||||
env_vec.push(format!(
|
||||
"http_proxy=http://{}:{}",
|
||||
proxy_host, self.proxy_port
|
||||
));
|
||||
env_vec.push(format!(
|
||||
"https_proxy=http://{}:{}",
|
||||
proxy_host, self.proxy_port
|
||||
));
|
||||
env_vec.push(format!(
|
||||
"HTTP_PROXY=http://{}:{}",
|
||||
proxy_host, self.proxy_port
|
||||
));
|
||||
env_vec.push(format!(
|
||||
"HTTPS_PROXY=http://{}:{}",
|
||||
proxy_host, self.proxy_port
|
||||
));
|
||||
}
|
||||
|
||||
// Build volume mounts based on policy
|
||||
let binds = match policy {
|
||||
SandboxPolicy::ReadOnly => {
|
||||
vec![format!("{}:/workspace:ro", working_dir_str)]
|
||||
}
|
||||
SandboxPolicy::WorkspaceWrite => {
|
||||
vec![format!("{}:/workspace:rw", working_dir_str)]
|
||||
}
|
||||
SandboxPolicy::FullAccess => {
|
||||
// Full access - mount more of the host
|
||||
vec![
|
||||
format!("{}:/workspace:rw", working_dir_str),
|
||||
"/tmp:/tmp:rw".to_string(),
|
||||
]
|
||||
}
|
||||
};
|
||||
|
||||
let host_config = HostConfig {
|
||||
binds: Some(binds),
|
||||
memory: Some((limits.memory_bytes) as i64),
|
||||
cpu_shares: Some(limits.cpu_shares as i64),
|
||||
auto_remove: Some(true),
|
||||
network_mode: Some("bridge".to_string()),
|
||||
// Security: drop all capabilities and add back only what's needed
|
||||
cap_drop: Some(vec!["ALL".to_string()]),
|
||||
cap_add: Some(vec![
|
||||
"CHOWN".to_string(),
|
||||
"SETUID".to_string(),
|
||||
"SETGID".to_string(),
|
||||
]),
|
||||
// Prevent privilege escalation
|
||||
security_opt: Some(vec!["no-new-privileges:true".to_string()]),
|
||||
// Read-only root filesystem (workspace is still writable if policy allows)
|
||||
readonly_rootfs: Some(policy == SandboxPolicy::ReadOnly),
|
||||
// Tmpfs mounts for /tmp and cargo cache
|
||||
tmpfs: Some(
|
||||
[
|
||||
("/tmp".to_string(), "size=512M".to_string()),
|
||||
(
|
||||
"/home/sandbox/.cargo/registry".to_string(),
|
||||
"size=1G".to_string(),
|
||||
),
|
||||
]
|
||||
.into_iter()
|
||||
.collect(),
|
||||
),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let config = Config {
|
||||
image: Some(self.image.clone()),
|
||||
cmd: Some(vec![
|
||||
"sh".to_string(),
|
||||
"-c".to_string(),
|
||||
command.to_string(),
|
||||
]),
|
||||
working_dir: Some("/workspace".to_string()),
|
||||
env: Some(env_vec),
|
||||
host_config: Some(host_config),
|
||||
user: Some("1000:1000".to_string()), // Non-root user
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let options = CreateContainerOptions {
|
||||
name: format!("sandbox-{}", uuid::Uuid::new_v4()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let response = self
|
||||
.docker
|
||||
.create_container(Some(options), config)
|
||||
.await
|
||||
.map_err(|e| SandboxError::ContainerCreationFailed {
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
Ok(response.id)
|
||||
}
|
||||
|
||||
/// Wait for a container to complete and collect output.
|
||||
async fn wait_for_container(
|
||||
&self,
|
||||
container_id: &str,
|
||||
max_output: usize,
|
||||
) -> Result<ContainerOutput> {
|
||||
// Wait for the container to finish
|
||||
let mut wait_stream = self.docker.wait_container(
|
||||
container_id,
|
||||
Some(WaitContainerOptions {
|
||||
condition: "not-running",
|
||||
}),
|
||||
);
|
||||
|
||||
let exit_code = match wait_stream.next().await {
|
||||
Some(Ok(response)) => response.status_code,
|
||||
Some(Err(e)) => {
|
||||
return Err(SandboxError::ExecutionFailed {
|
||||
reason: format!("wait failed: {}", e),
|
||||
});
|
||||
}
|
||||
None => {
|
||||
return Err(SandboxError::ExecutionFailed {
|
||||
reason: "container wait stream ended unexpectedly".to_string(),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
// Collect logs
|
||||
let (stdout, stderr, truncated) = self.collect_logs(container_id, max_output).await?;
|
||||
|
||||
Ok(ContainerOutput {
|
||||
exit_code,
|
||||
stdout,
|
||||
stderr,
|
||||
duration: Duration::ZERO, // Will be set by caller
|
||||
truncated,
|
||||
})
|
||||
}
|
||||
|
||||
/// Collect stdout and stderr from a container.
|
||||
async fn collect_logs(
|
||||
&self,
|
||||
container_id: &str,
|
||||
max_output: usize,
|
||||
) -> Result<(String, String, bool)> {
|
||||
let options = LogsOptions::<String> {
|
||||
stdout: true,
|
||||
stderr: true,
|
||||
follow: false,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let mut stream = self.docker.logs(container_id, Some(options));
|
||||
|
||||
let mut stdout = String::new();
|
||||
let mut stderr = String::new();
|
||||
let mut truncated = false;
|
||||
let half_max = max_output / 2;
|
||||
|
||||
while let Some(result) = stream.next().await {
|
||||
match result {
|
||||
Ok(LogOutput::StdOut { message }) => {
|
||||
let text = String::from_utf8_lossy(&message);
|
||||
if stdout.len() + text.len() > half_max {
|
||||
truncated = true;
|
||||
let remaining = half_max.saturating_sub(stdout.len());
|
||||
stdout.push_str(&text[..remaining.min(text.len())]);
|
||||
} else {
|
||||
stdout.push_str(&text);
|
||||
}
|
||||
}
|
||||
Ok(LogOutput::StdErr { message }) => {
|
||||
let text = String::from_utf8_lossy(&message);
|
||||
if stderr.len() + text.len() > half_max {
|
||||
truncated = true;
|
||||
let remaining = half_max.saturating_sub(stderr.len());
|
||||
stderr.push_str(&text[..remaining.min(text.len())]);
|
||||
} else {
|
||||
stderr.push_str(&text);
|
||||
}
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
tracing::warn!("Error reading container logs: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok((stdout, stderr, truncated))
|
||||
}
|
||||
|
||||
/// Run an exec and collect output.
|
||||
async fn run_exec(&self, exec_id: &str, max_output: usize) -> Result<ContainerOutput> {
|
||||
let start_result = self.docker.start_exec(exec_id, None).await.map_err(|e| {
|
||||
SandboxError::ExecutionFailed {
|
||||
reason: format!("exec start failed: {}", e),
|
||||
}
|
||||
})?;
|
||||
|
||||
let mut stdout = String::new();
|
||||
let mut stderr = String::new();
|
||||
let mut truncated = false;
|
||||
let half_max = max_output / 2;
|
||||
|
||||
if let StartExecResults::Attached { mut output, .. } = start_result {
|
||||
while let Some(result) = output.next().await {
|
||||
match result {
|
||||
Ok(LogOutput::StdOut { message }) => {
|
||||
let text = String::from_utf8_lossy(&message);
|
||||
if stdout.len() < half_max {
|
||||
let remaining = half_max.saturating_sub(stdout.len());
|
||||
stdout.push_str(&text[..remaining.min(text.len())]);
|
||||
if text.len() > remaining {
|
||||
truncated = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(LogOutput::StdErr { message }) => {
|
||||
let text = String::from_utf8_lossy(&message);
|
||||
if stderr.len() < half_max {
|
||||
let remaining = half_max.saturating_sub(stderr.len());
|
||||
stderr.push_str(&text[..remaining.min(text.len())]);
|
||||
if text.len() > remaining {
|
||||
truncated = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
tracing::warn!("Error reading exec output: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Get exec exit code
|
||||
let inspect =
|
||||
self.docker
|
||||
.inspect_exec(exec_id)
|
||||
.await
|
||||
.map_err(|e| SandboxError::ExecutionFailed {
|
||||
reason: format!("exec inspect failed: {}", e),
|
||||
})?;
|
||||
|
||||
let exit_code = inspect.exit_code.unwrap_or(-1);
|
||||
|
||||
Ok(ContainerOutput {
|
||||
exit_code,
|
||||
stdout,
|
||||
stderr,
|
||||
duration: Duration::ZERO,
|
||||
truncated,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Connect to the Docker daemon.
|
||||
pub async fn connect_docker() -> Result<Docker> {
|
||||
Docker::connect_with_local_defaults().map_err(|e| SandboxError::DockerNotAvailable {
|
||||
reason: e.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_docker_connection() {
|
||||
// This test requires Docker to be running
|
||||
let result = connect_docker().await;
|
||||
// Don't fail if Docker isn't available, just skip
|
||||
if result.is_err() {
|
||||
eprintln!("Skipping Docker test: Docker not available");
|
||||
return;
|
||||
}
|
||||
|
||||
let docker = result.unwrap();
|
||||
let runner = ContainerRunner::new(docker, "alpine:latest".to_string(), 0);
|
||||
// Just check that we can query Docker (result doesn't matter for CI)
|
||||
let _available = runner.is_available().await;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
//! Error types for the Docker execution sandbox.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
/// Errors that can occur in the sandbox system.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum SandboxError {
|
||||
/// Docker daemon is not available or not running.
|
||||
#[error("Docker not available: {reason}")]
|
||||
DockerNotAvailable { reason: String },
|
||||
|
||||
/// Failed to create container.
|
||||
#[error("Container creation failed: {reason}")]
|
||||
ContainerCreationFailed { reason: String },
|
||||
|
||||
/// Failed to start container.
|
||||
#[error("Container start failed: {reason}")]
|
||||
ContainerStartFailed { reason: String },
|
||||
|
||||
/// Command execution failed inside container.
|
||||
#[error("Execution failed: {reason}")]
|
||||
ExecutionFailed { reason: String },
|
||||
|
||||
/// Command timed out.
|
||||
#[error("Command timed out after {0:?}")]
|
||||
Timeout(Duration),
|
||||
|
||||
/// Container resource limit exceeded.
|
||||
#[error("Resource limit exceeded: {resource} limit of {limit}")]
|
||||
ResourceLimitExceeded { resource: String, limit: String },
|
||||
|
||||
/// Network proxy error.
|
||||
#[error("Proxy error: {reason}")]
|
||||
ProxyError { reason: String },
|
||||
|
||||
/// Network request blocked by policy.
|
||||
#[error("Network request blocked: {reason}")]
|
||||
NetworkBlocked { reason: String },
|
||||
|
||||
/// Credential injection failed.
|
||||
#[error("Credential injection failed for {domain}: {reason}")]
|
||||
CredentialInjectionFailed { domain: String, reason: String },
|
||||
|
||||
/// Docker API error.
|
||||
#[error("Docker API error: {0}")]
|
||||
Docker(#[from] bollard::errors::Error),
|
||||
|
||||
/// I/O error.
|
||||
#[error("I/O error: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
|
||||
/// Configuration error.
|
||||
#[error("Configuration error: {reason}")]
|
||||
Config { reason: String },
|
||||
}
|
||||
|
||||
/// Result type for sandbox operations.
|
||||
pub type Result<T> = std::result::Result<T, SandboxError>;
|
||||
@@ -0,0 +1,474 @@
|
||||
//! Main sandbox manager coordinating proxy and containers.
|
||||
//!
|
||||
//! The `SandboxManager` is the primary entry point for sandboxed execution.
|
||||
//! It coordinates:
|
||||
//! - Docker container creation and lifecycle
|
||||
//! - HTTP proxy for network access control
|
||||
//! - Credential injection for API calls
|
||||
//! - Resource limits and timeouts
|
||||
//!
|
||||
//! # Architecture
|
||||
//!
|
||||
//! ```text
|
||||
//! ┌───────────────────────────────────────────────────────────────────────────┐
|
||||
//! │ SandboxManager │
|
||||
//! │ │
|
||||
//! │ execute(cmd, cwd, policy) │
|
||||
//! │ │ │
|
||||
//! │ ▼ │
|
||||
//! │ ┌──────────────┐ ┌──────────────┐ ┌──────────────────────────┐ │
|
||||
//! │ │ Start Proxy │────▶│ Create │────▶│ Execute & Collect Output │ │
|
||||
//! │ │ (if needed) │ │ Container │ │ │ │
|
||||
//! │ └──────────────┘ └──────────────┘ └──────────────────────────┘ │
|
||||
//! │ │ │
|
||||
//! │ ▼ │
|
||||
//! │ ┌──────────────────────────┐ │
|
||||
//! │ │ Cleanup Container │ │
|
||||
//! │ └──────────────────────────┘ │
|
||||
//! └───────────────────────────────────────────────────────────────────────────┘
|
||||
//! ```
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::sandbox::config::{ResourceLimits, SandboxConfig, SandboxPolicy};
|
||||
use crate::sandbox::container::{ContainerOutput, ContainerRunner, connect_docker};
|
||||
use crate::sandbox::error::{Result, SandboxError};
|
||||
use crate::sandbox::proxy::{HttpProxy, NetworkProxyBuilder};
|
||||
|
||||
/// Output from sandbox execution.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ExecOutput {
|
||||
/// Exit code from the command.
|
||||
pub exit_code: i64,
|
||||
/// Standard output.
|
||||
pub stdout: String,
|
||||
/// Standard error.
|
||||
pub stderr: String,
|
||||
/// Combined output (stdout + stderr).
|
||||
pub output: String,
|
||||
/// How long the command ran.
|
||||
pub duration: Duration,
|
||||
/// Whether output was truncated.
|
||||
pub truncated: bool,
|
||||
}
|
||||
|
||||
impl From<ContainerOutput> for ExecOutput {
|
||||
fn from(c: ContainerOutput) -> Self {
|
||||
let output = if c.stderr.is_empty() {
|
||||
c.stdout.clone()
|
||||
} else if c.stdout.is_empty() {
|
||||
c.stderr.clone()
|
||||
} else {
|
||||
format!("{}\n\n--- stderr ---\n{}", c.stdout, c.stderr)
|
||||
};
|
||||
|
||||
Self {
|
||||
exit_code: c.exit_code,
|
||||
stdout: c.stdout,
|
||||
stderr: c.stderr,
|
||||
output,
|
||||
duration: c.duration,
|
||||
truncated: c.truncated,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Main sandbox manager.
|
||||
pub struct SandboxManager {
|
||||
config: SandboxConfig,
|
||||
proxy: Arc<RwLock<Option<HttpProxy>>>,
|
||||
runner: Arc<RwLock<Option<ContainerRunner>>>,
|
||||
initialized: std::sync::atomic::AtomicBool,
|
||||
}
|
||||
|
||||
impl SandboxManager {
|
||||
/// Create a new sandbox manager.
|
||||
pub fn new(config: SandboxConfig) -> Self {
|
||||
Self {
|
||||
config,
|
||||
proxy: Arc::new(RwLock::new(None)),
|
||||
runner: Arc::new(RwLock::new(None)),
|
||||
initialized: std::sync::atomic::AtomicBool::new(false),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create with default configuration.
|
||||
pub fn with_defaults() -> Self {
|
||||
Self::new(SandboxConfig::default())
|
||||
}
|
||||
|
||||
/// Check if the sandbox is available (Docker running, etc.).
|
||||
pub async fn is_available(&self) -> bool {
|
||||
if !self.config.enabled {
|
||||
return false;
|
||||
}
|
||||
|
||||
match connect_docker().await {
|
||||
Ok(docker) => docker.ping().await.is_ok(),
|
||||
Err(_) => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Initialize the sandbox (connect to Docker, start proxy).
|
||||
pub async fn initialize(&self) -> Result<()> {
|
||||
if self.initialized.load(std::sync::atomic::Ordering::SeqCst) {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if !self.config.enabled {
|
||||
return Err(SandboxError::Config {
|
||||
reason: "sandbox is disabled".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Connect to Docker
|
||||
let docker = connect_docker().await?;
|
||||
|
||||
// Check if Docker is responsive
|
||||
docker
|
||||
.ping()
|
||||
.await
|
||||
.map_err(|e| SandboxError::DockerNotAvailable {
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
// Create container runner
|
||||
let runner =
|
||||
ContainerRunner::new(docker, self.config.image.clone(), self.config.proxy_port);
|
||||
|
||||
// Check for / pull image
|
||||
if !runner.image_exists().await {
|
||||
if self.config.auto_pull_image {
|
||||
runner.pull_image().await?;
|
||||
} else {
|
||||
return Err(SandboxError::ContainerCreationFailed {
|
||||
reason: format!(
|
||||
"image {} not found and auto_pull is disabled",
|
||||
self.config.image
|
||||
),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
*self.runner.write().await = Some(runner);
|
||||
|
||||
// Start the network proxy if we're using a sandboxed policy
|
||||
if self.config.policy.is_sandboxed() {
|
||||
let proxy = NetworkProxyBuilder::from_config(&self.config)
|
||||
.build_and_start(self.config.proxy_port)
|
||||
.await?;
|
||||
|
||||
*self.proxy.write().await = Some(proxy);
|
||||
}
|
||||
|
||||
self.initialized
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
|
||||
tracing::info!("Sandbox initialized");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Shutdown the sandbox (stop proxy, clean up).
|
||||
pub async fn shutdown(&self) {
|
||||
if let Some(proxy) = self.proxy.write().await.take() {
|
||||
proxy.stop().await;
|
||||
}
|
||||
|
||||
self.initialized
|
||||
.store(false, std::sync::atomic::Ordering::SeqCst);
|
||||
|
||||
tracing::info!("Sandbox shut down");
|
||||
}
|
||||
|
||||
/// Execute a command in the sandbox.
|
||||
pub async fn execute(
|
||||
&self,
|
||||
command: &str,
|
||||
cwd: &Path,
|
||||
env: HashMap<String, String>,
|
||||
) -> Result<ExecOutput> {
|
||||
self.execute_with_policy(command, cwd, self.config.policy, env)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Execute a command with a specific policy.
|
||||
pub async fn execute_with_policy(
|
||||
&self,
|
||||
command: &str,
|
||||
cwd: &Path,
|
||||
policy: SandboxPolicy,
|
||||
env: HashMap<String, String>,
|
||||
) -> Result<ExecOutput> {
|
||||
// FullAccess policy bypasses the sandbox entirely
|
||||
if policy == SandboxPolicy::FullAccess {
|
||||
return self.execute_direct(command, cwd, env).await;
|
||||
}
|
||||
|
||||
// Ensure we're initialized
|
||||
if !self.initialized.load(std::sync::atomic::Ordering::SeqCst) {
|
||||
self.initialize().await?;
|
||||
}
|
||||
|
||||
// Get proxy port if running
|
||||
let proxy_port = if let Some(proxy) = self.proxy.read().await.as_ref() {
|
||||
proxy.addr().await.map(|a| a.port()).unwrap_or(0)
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
// Create a runner with the current proxy port
|
||||
let docker = connect_docker().await?;
|
||||
let runner = ContainerRunner::new(docker, self.config.image.clone(), proxy_port);
|
||||
|
||||
let limits = ResourceLimits {
|
||||
memory_bytes: self.config.memory_limit_mb * 1024 * 1024,
|
||||
cpu_shares: self.config.cpu_shares,
|
||||
timeout: self.config.timeout,
|
||||
max_output_bytes: 64 * 1024,
|
||||
};
|
||||
|
||||
let container_output = runner.execute(command, cwd, policy, &limits, env).await?;
|
||||
|
||||
Ok(container_output.into())
|
||||
}
|
||||
|
||||
/// Execute a command directly on the host (no sandbox).
|
||||
async fn execute_direct(
|
||||
&self,
|
||||
command: &str,
|
||||
cwd: &Path,
|
||||
env: HashMap<String, String>,
|
||||
) -> Result<ExecOutput> {
|
||||
use tokio::process::Command;
|
||||
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
let mut cmd = if cfg!(target_os = "windows") {
|
||||
let mut c = Command::new("cmd");
|
||||
c.args(["/C", command]);
|
||||
c
|
||||
} else {
|
||||
let mut c = Command::new("sh");
|
||||
c.args(["-c", command]);
|
||||
c
|
||||
};
|
||||
|
||||
cmd.current_dir(cwd);
|
||||
cmd.envs(env);
|
||||
|
||||
let output = tokio::time::timeout(self.config.timeout, cmd.output())
|
||||
.await
|
||||
.map_err(|_| SandboxError::Timeout(self.config.timeout))?
|
||||
.map_err(|e| SandboxError::ExecutionFailed {
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
let stdout = String::from_utf8_lossy(&output.stdout).to_string();
|
||||
let stderr = String::from_utf8_lossy(&output.stderr).to_string();
|
||||
let combined = if stderr.is_empty() {
|
||||
stdout.clone()
|
||||
} else if stdout.is_empty() {
|
||||
stderr.clone()
|
||||
} else {
|
||||
format!("{}\n\n--- stderr ---\n{}", stdout, stderr)
|
||||
};
|
||||
|
||||
Ok(ExecOutput {
|
||||
exit_code: output.status.code().unwrap_or(-1) as i64,
|
||||
stdout,
|
||||
stderr,
|
||||
output: combined,
|
||||
duration: start.elapsed(),
|
||||
truncated: false,
|
||||
})
|
||||
}
|
||||
|
||||
/// Execute a build command (convenience method using WorkspaceWrite policy).
|
||||
pub async fn build(
|
||||
&self,
|
||||
command: &str,
|
||||
project_dir: &Path,
|
||||
env: HashMap<String, String>,
|
||||
) -> Result<ExecOutput> {
|
||||
self.execute_with_policy(command, project_dir, SandboxPolicy::WorkspaceWrite, env)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Get the current configuration.
|
||||
pub fn config(&self) -> &SandboxConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// Check if the sandbox is initialized.
|
||||
pub fn is_initialized(&self) -> bool {
|
||||
self.initialized.load(std::sync::atomic::Ordering::SeqCst)
|
||||
}
|
||||
|
||||
/// Get the proxy port if running.
|
||||
pub async fn proxy_port(&self) -> Option<u16> {
|
||||
if let Some(proxy) = self.proxy.read().await.as_ref() {
|
||||
proxy.addr().await.map(|a| a.port())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for SandboxManager {
|
||||
fn drop(&mut self) {
|
||||
// Note: async cleanup should be done via shutdown() before dropping
|
||||
if self.initialized.load(std::sync::atomic::Ordering::SeqCst) {
|
||||
tracing::warn!("SandboxManager dropped without shutdown(), resources may leak");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Builder for creating a sandbox manager.
|
||||
pub struct SandboxManagerBuilder {
|
||||
config: SandboxConfig,
|
||||
}
|
||||
|
||||
impl SandboxManagerBuilder {
|
||||
/// Create a new builder.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
config: SandboxConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Enable the sandbox.
|
||||
pub fn enabled(mut self, enabled: bool) -> Self {
|
||||
self.config.enabled = enabled;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the sandbox policy.
|
||||
pub fn policy(mut self, policy: SandboxPolicy) -> Self {
|
||||
self.config.policy = policy;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the command timeout.
|
||||
pub fn timeout(mut self, timeout: Duration) -> Self {
|
||||
self.config.timeout = timeout;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the memory limit in MB.
|
||||
pub fn memory_limit_mb(mut self, mb: u64) -> Self {
|
||||
self.config.memory_limit_mb = mb;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the Docker image.
|
||||
pub fn image(mut self, image: &str) -> Self {
|
||||
self.config.image = image.to_string();
|
||||
self
|
||||
}
|
||||
|
||||
/// Add domains to the network allowlist.
|
||||
pub fn allow_domains(mut self, domains: Vec<String>) -> Self {
|
||||
self.config.network_allowlist.extend(domains);
|
||||
self
|
||||
}
|
||||
|
||||
/// Build the sandbox manager.
|
||||
pub fn build(self) -> SandboxManager {
|
||||
SandboxManager::new(self.config)
|
||||
}
|
||||
|
||||
/// Build and initialize the sandbox manager.
|
||||
pub async fn build_and_init(self) -> Result<SandboxManager> {
|
||||
let manager = self.build();
|
||||
manager.initialize().await?;
|
||||
Ok(manager)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for SandboxManagerBuilder {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_exec_output_from_container_output() {
|
||||
let container = ContainerOutput {
|
||||
exit_code: 0,
|
||||
stdout: "hello".to_string(),
|
||||
stderr: String::new(),
|
||||
duration: Duration::from_secs(1),
|
||||
truncated: false,
|
||||
};
|
||||
|
||||
let exec: ExecOutput = container.into();
|
||||
assert_eq!(exec.exit_code, 0);
|
||||
assert_eq!(exec.output, "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exec_output_combined() {
|
||||
let container = ContainerOutput {
|
||||
exit_code: 1,
|
||||
stdout: "out".to_string(),
|
||||
stderr: "err".to_string(),
|
||||
duration: Duration::from_secs(1),
|
||||
truncated: false,
|
||||
};
|
||||
|
||||
let exec: ExecOutput = container.into();
|
||||
assert!(exec.output.contains("out"));
|
||||
assert!(exec.output.contains("err"));
|
||||
assert!(exec.output.contains("stderr"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_builder_defaults() {
|
||||
let manager = SandboxManagerBuilder::new().build();
|
||||
assert!(!manager.config.enabled); // Disabled by default
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_builder_custom() {
|
||||
let manager = SandboxManagerBuilder::new()
|
||||
.enabled(true)
|
||||
.policy(SandboxPolicy::WorkspaceWrite)
|
||||
.timeout(Duration::from_secs(60))
|
||||
.memory_limit_mb(1024)
|
||||
.image("custom:latest")
|
||||
.build();
|
||||
|
||||
assert!(manager.config.enabled);
|
||||
assert_eq!(manager.config.policy, SandboxPolicy::WorkspaceWrite);
|
||||
assert_eq!(manager.config.timeout, Duration::from_secs(60));
|
||||
assert_eq!(manager.config.memory_limit_mb, 1024);
|
||||
assert_eq!(manager.config.image, "custom:latest");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_direct_execution() {
|
||||
let manager = SandboxManager::new(SandboxConfig {
|
||||
enabled: true,
|
||||
policy: SandboxPolicy::FullAccess,
|
||||
..Default::default()
|
||||
});
|
||||
|
||||
let result = manager
|
||||
.execute("echo hello", Path::new("."), HashMap::new())
|
||||
.await;
|
||||
|
||||
// This should work even without Docker since FullAccess runs directly
|
||||
assert!(result.is_ok());
|
||||
let output = result.unwrap();
|
||||
assert!(output.stdout.contains("hello"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
//! Docker execution sandbox for secure command execution.
|
||||
//!
|
||||
//! This module provides a complete sandboxing solution for running untrusted commands:
|
||||
//! - **Container isolation**: Commands run in ephemeral Docker containers
|
||||
//! - **Network proxy**: All network traffic goes through a validating proxy
|
||||
//! - **Credential injection**: Secrets are injected by the proxy, never exposed in containers
|
||||
//! - **Resource limits**: Memory, CPU, and timeout enforcement
|
||||
//!
|
||||
//! # Architecture
|
||||
//!
|
||||
//! ```text
|
||||
//! ┌─────────────────────────────────────────────────────────────────────────────┐
|
||||
//! │ Sandbox System │
|
||||
//! │ │
|
||||
//! │ ┌─────────────────────────────────────────────────────────────────────┐ │
|
||||
//! │ │ SandboxManager │ │
|
||||
//! │ │ │ │
|
||||
//! │ │ • Coordinates container creation and execution │ │
|
||||
//! │ │ • Manages proxy lifecycle │ │
|
||||
//! │ │ • Enforces resource limits │ │
|
||||
//! │ └─────────────────────────────────────────────────────────────────────┘ │
|
||||
//! │ │ │ │
|
||||
//! │ ▼ ▼ │
|
||||
//! │ ┌──────────────────┐ ┌───────────────────┐ │
|
||||
//! │ │ Container │ │ Network Proxy │ │
|
||||
//! │ │ Runner │ │ │ │
|
||||
//! │ │ │ │ • Allowlist │ │
|
||||
//! │ │ • Create │◀────────▶│ • Credentials │ │
|
||||
//! │ │ • Execute │ │ • Logging │ │
|
||||
//! │ │ • Cleanup │ │ │ │
|
||||
//! │ └──────────────────┘ └───────────────────┘ │
|
||||
//! │ │ │ │
|
||||
//! │ ▼ ▼ │
|
||||
//! │ ┌──────────────────┐ ┌───────────────────┐ │
|
||||
//! │ │ Docker │ │ Internet │ │
|
||||
//! │ │ │ │ (allowed hosts) │ │
|
||||
//! │ └──────────────────┘ └───────────────────┘ │
|
||||
//! └─────────────────────────────────────────────────────────────────────────────┘
|
||||
//! ```
|
||||
//!
|
||||
//! # Sandbox Policies
|
||||
//!
|
||||
//! | Policy | Filesystem | Network | Use Case |
|
||||
//! |--------|------------|---------|----------|
|
||||
//! | `ReadOnly` | Read workspace | Proxied | Explore code, fetch docs |
|
||||
//! | `WorkspaceWrite` | Read/write workspace | Proxied | Build software, run tests |
|
||||
//! | `FullAccess` | Full host | Full | Direct execution (no sandbox) |
|
||||
//!
|
||||
//! # Example
|
||||
//!
|
||||
//! ```rust,no_run
|
||||
//! use near_agent::sandbox::{SandboxManager, SandboxManagerBuilder, SandboxPolicy};
|
||||
//! use std::collections::HashMap;
|
||||
//! use std::path::Path;
|
||||
//!
|
||||
//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
|
||||
//! let manager = SandboxManagerBuilder::new()
|
||||
//! .enabled(true)
|
||||
//! .policy(SandboxPolicy::WorkspaceWrite)
|
||||
//! .build();
|
||||
//!
|
||||
//! manager.initialize().await?;
|
||||
//!
|
||||
//! let result = manager.execute(
|
||||
//! "cargo build --release",
|
||||
//! Path::new("/workspace/my-project"),
|
||||
//! HashMap::new(),
|
||||
//! ).await?;
|
||||
//!
|
||||
//! println!("Exit code: {}", result.exit_code);
|
||||
//! println!("Output: {}", result.output);
|
||||
//!
|
||||
//! manager.shutdown().await;
|
||||
//! # Ok(())
|
||||
//! # }
|
||||
//! ```
|
||||
//!
|
||||
//! # Security Properties
|
||||
//!
|
||||
//! - **No credentials in containers**: Environment variables with secrets never enter containers
|
||||
//! - **Network isolation**: All traffic routes through the proxy (validated domains only)
|
||||
//! - **Non-root execution**: Containers run as UID 1000
|
||||
//! - **Read-only root**: Container filesystem is read-only (except workspace mount)
|
||||
//! - **Capability dropping**: All Linux capabilities dropped, only essential ones added back
|
||||
//! - **Auto-cleanup**: Containers are removed after execution (--rm + explicit cleanup)
|
||||
//! - **Timeout enforcement**: Commands are killed after the timeout
|
||||
|
||||
pub mod config;
|
||||
pub mod container;
|
||||
pub mod error;
|
||||
pub mod manager;
|
||||
pub mod proxy;
|
||||
|
||||
pub use config::{
|
||||
CredentialLocation, CredentialMapping, ResourceLimits, SandboxConfig, SandboxPolicy,
|
||||
};
|
||||
pub use container::{ContainerOutput, ContainerRunner, connect_docker};
|
||||
pub use error::{Result, SandboxError};
|
||||
pub use manager::{ExecOutput, SandboxManager, SandboxManagerBuilder};
|
||||
pub use proxy::{
|
||||
CredentialResolver, DefaultPolicyDecider, DomainAllowlist, EnvCredentialResolver, HttpProxy,
|
||||
NetworkDecision, NetworkPolicyDecider, NetworkProxyBuilder, NetworkRequest,
|
||||
};
|
||||
|
||||
/// Default allowlist getter (re-export for convenience).
|
||||
pub fn default_allowlist() -> Vec<String> {
|
||||
config::default_allowlist()
|
||||
}
|
||||
|
||||
/// Default credential mappings getter (re-export for convenience).
|
||||
pub fn default_credential_mappings() -> Vec<CredentialMapping> {
|
||||
config::default_credential_mappings()
|
||||
}
|
||||
@@ -0,0 +1,251 @@
|
||||
//! Domain allowlist for the network proxy.
|
||||
//!
|
||||
//! Validates that HTTP requests only go to allowed domains.
|
||||
//! Supports exact matches and wildcard patterns.
|
||||
|
||||
use std::fmt;
|
||||
|
||||
/// Pattern for matching allowed domains.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DomainPattern {
|
||||
/// The domain pattern (e.g., "api.example.com" or "*.example.com").
|
||||
pattern: String,
|
||||
/// Whether this is a wildcard pattern.
|
||||
is_wildcard: bool,
|
||||
/// The base domain for wildcard matching.
|
||||
base_domain: String,
|
||||
}
|
||||
|
||||
impl DomainPattern {
|
||||
/// Create a new domain pattern.
|
||||
pub fn new(pattern: &str) -> Self {
|
||||
let is_wildcard = pattern.starts_with("*.");
|
||||
let base_domain = if is_wildcard {
|
||||
pattern[2..].to_lowercase()
|
||||
} else {
|
||||
pattern.to_lowercase()
|
||||
};
|
||||
|
||||
Self {
|
||||
pattern: pattern.to_string(),
|
||||
is_wildcard,
|
||||
base_domain,
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if a host matches this pattern.
|
||||
pub fn matches(&self, host: &str) -> bool {
|
||||
let host_lower = host.to_lowercase();
|
||||
|
||||
if self.is_wildcard {
|
||||
// *.example.com matches foo.example.com, bar.baz.example.com, example.com
|
||||
host_lower == self.base_domain
|
||||
|| host_lower.ends_with(&format!(".{}", self.base_domain))
|
||||
} else {
|
||||
host_lower == self.base_domain
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the pattern string.
|
||||
pub fn pattern(&self) -> &str {
|
||||
&self.pattern
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for DomainPattern {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{}", self.pattern)
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of domain validation.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum DomainValidationResult {
|
||||
/// Domain is allowed.
|
||||
Allowed,
|
||||
/// Domain is denied with a reason.
|
||||
Denied(String),
|
||||
}
|
||||
|
||||
impl DomainValidationResult {
|
||||
pub fn is_allowed(&self) -> bool {
|
||||
matches!(self, DomainValidationResult::Allowed)
|
||||
}
|
||||
}
|
||||
|
||||
/// Validates domains against an allowlist.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DomainAllowlist {
|
||||
patterns: Vec<DomainPattern>,
|
||||
}
|
||||
|
||||
impl DomainAllowlist {
|
||||
/// Create a new allowlist from domain strings.
|
||||
pub fn new(domains: &[String]) -> Self {
|
||||
Self {
|
||||
patterns: domains.iter().map(|d| DomainPattern::new(d)).collect(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an empty allowlist (denies everything).
|
||||
pub fn empty() -> Self {
|
||||
Self { patterns: vec![] }
|
||||
}
|
||||
|
||||
/// Add a domain pattern to the allowlist.
|
||||
pub fn add(&mut self, pattern: &str) {
|
||||
self.patterns.push(DomainPattern::new(pattern));
|
||||
}
|
||||
|
||||
/// Check if a domain is allowed.
|
||||
pub fn is_allowed(&self, host: &str) -> DomainValidationResult {
|
||||
if self.patterns.is_empty() {
|
||||
return DomainValidationResult::Denied("empty allowlist".to_string());
|
||||
}
|
||||
|
||||
for pattern in &self.patterns {
|
||||
if pattern.matches(host) {
|
||||
return DomainValidationResult::Allowed;
|
||||
}
|
||||
}
|
||||
|
||||
DomainValidationResult::Denied(format!(
|
||||
"host '{}' not in allowlist: [{}]",
|
||||
host,
|
||||
self.patterns
|
||||
.iter()
|
||||
.map(|p| p.pattern())
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ")
|
||||
))
|
||||
}
|
||||
|
||||
/// Get all patterns in the allowlist.
|
||||
pub fn patterns(&self) -> &[DomainPattern] {
|
||||
&self.patterns
|
||||
}
|
||||
|
||||
/// Check if the allowlist is empty.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.patterns.is_empty()
|
||||
}
|
||||
|
||||
/// Get the number of patterns.
|
||||
pub fn len(&self) -> usize {
|
||||
self.patterns.len()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for DomainAllowlist {
|
||||
fn default() -> Self {
|
||||
Self::new(&crate::sandbox::config::default_allowlist())
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse host from a URL string.
|
||||
pub fn extract_host(url: &str) -> Option<String> {
|
||||
// Determine scheme and extract the rest
|
||||
let rest = if let Some(stripped) = url.strip_prefix("https://") {
|
||||
stripped
|
||||
} else if let Some(stripped) = url.strip_prefix("http://") {
|
||||
stripped
|
||||
} else {
|
||||
return None;
|
||||
};
|
||||
|
||||
// Find the end of the host (start of path, query, or end of string)
|
||||
let host_end = rest.find('/').unwrap_or(rest.len());
|
||||
let host_and_port = &rest[..host_end];
|
||||
|
||||
// Remove port if present
|
||||
let host = if let Some(bracket_idx) = host_and_port.find('[') {
|
||||
// IPv6 address
|
||||
let close_bracket = host_and_port.find(']')?;
|
||||
&host_and_port[bracket_idx + 1..close_bracket]
|
||||
} else if let Some(colon_idx) = host_and_port.rfind(':') {
|
||||
// Check if this is a port (all digits after colon)
|
||||
let after_colon = &host_and_port[colon_idx + 1..];
|
||||
if after_colon.chars().all(|c| c.is_ascii_digit()) {
|
||||
&host_and_port[..colon_idx]
|
||||
} else {
|
||||
host_and_port
|
||||
}
|
||||
} else {
|
||||
host_and_port
|
||||
};
|
||||
|
||||
if host.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(host.to_lowercase())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_exact_match() {
|
||||
let pattern = DomainPattern::new("api.example.com");
|
||||
assert!(pattern.matches("api.example.com"));
|
||||
assert!(pattern.matches("API.EXAMPLE.COM"));
|
||||
assert!(!pattern.matches("foo.api.example.com"));
|
||||
assert!(!pattern.matches("example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wildcard_match() {
|
||||
let pattern = DomainPattern::new("*.example.com");
|
||||
assert!(pattern.matches("api.example.com"));
|
||||
assert!(pattern.matches("foo.bar.example.com"));
|
||||
assert!(pattern.matches("example.com")); // Base domain also matches
|
||||
assert!(!pattern.matches("exampleXcom"));
|
||||
assert!(!pattern.matches("other.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_allowlist_allows() {
|
||||
let allowlist =
|
||||
DomainAllowlist::new(&["crates.io".to_string(), "*.github.com".to_string()]);
|
||||
|
||||
assert!(allowlist.is_allowed("crates.io").is_allowed());
|
||||
assert!(allowlist.is_allowed("api.github.com").is_allowed());
|
||||
assert!(
|
||||
!allowlist
|
||||
.is_allowed("raw.githubusercontent.com")
|
||||
.is_allowed()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_allowlist_denies() {
|
||||
let allowlist = DomainAllowlist::new(&["crates.io".to_string()]);
|
||||
|
||||
let result = allowlist.is_allowed("evil.com");
|
||||
assert!(!result.is_allowed());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_allowlist() {
|
||||
let allowlist = DomainAllowlist::empty();
|
||||
assert!(!allowlist.is_allowed("anything.com").is_allowed());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_host() {
|
||||
assert_eq!(
|
||||
extract_host("https://api.example.com/v1/endpoint"),
|
||||
Some("api.example.com".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
extract_host("http://localhost:8080/api"),
|
||||
Some("localhost".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
extract_host("https://EXAMPLE.COM"),
|
||||
Some("example.com".to_string())
|
||||
);
|
||||
assert_eq!(extract_host("not-a-url"), None);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,444 @@
|
||||
//! HTTP proxy server for sandboxed network access.
|
||||
//!
|
||||
//! This proxy runs on the host and handles all network requests from containers.
|
||||
//! It validates requests against the allowlist and injects credentials when needed.
|
||||
//!
|
||||
//! ```text
|
||||
//! Container ──► http_proxy=host.docker.internal:PORT ──► This Proxy ──► Internet
|
||||
//! │
|
||||
//! ├─► Validate domain
|
||||
//! ├─► Inject credentials
|
||||
//! └─► Log requests
|
||||
//! ```
|
||||
|
||||
use std::convert::Infallible;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
|
||||
use bytes::Bytes;
|
||||
use http_body_util::{BodyExt, Empty, Full, combinators::BoxBody};
|
||||
use hyper::server::conn::http1;
|
||||
use hyper::service::service_fn;
|
||||
use hyper::{Method, Request, Response, StatusCode};
|
||||
use hyper_util::rt::TokioIo;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::sandbox::config::CredentialLocation;
|
||||
use crate::sandbox::error::{Result, SandboxError};
|
||||
use crate::sandbox::proxy::policy::{NetworkDecision, NetworkPolicyDecider, NetworkRequest};
|
||||
|
||||
/// State shared across proxy connections.
|
||||
struct ProxyState {
|
||||
/// Policy decider for network requests.
|
||||
decider: Arc<dyn NetworkPolicyDecider>,
|
||||
/// Credential resolver (maps secret names to values).
|
||||
credential_resolver: Arc<dyn CredentialResolver>,
|
||||
/// Request counter for logging.
|
||||
request_count: std::sync::atomic::AtomicU64,
|
||||
/// Whether the proxy is running.
|
||||
running: std::sync::atomic::AtomicBool,
|
||||
}
|
||||
|
||||
/// Resolves secret names to their values.
|
||||
#[async_trait::async_trait]
|
||||
pub trait CredentialResolver: Send + Sync {
|
||||
/// Get the value of a secret by name.
|
||||
async fn resolve(&self, name: &str) -> Option<String>;
|
||||
}
|
||||
|
||||
/// A credential resolver that uses environment variables.
|
||||
pub struct EnvCredentialResolver;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl CredentialResolver for EnvCredentialResolver {
|
||||
async fn resolve(&self, name: &str) -> Option<String> {
|
||||
std::env::var(name).ok()
|
||||
}
|
||||
}
|
||||
|
||||
/// A credential resolver that returns nothing (for testing).
|
||||
pub struct NoCredentialResolver;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl CredentialResolver for NoCredentialResolver {
|
||||
async fn resolve(&self, _name: &str) -> Option<String> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// HTTP proxy server.
|
||||
pub struct HttpProxy {
|
||||
state: Arc<ProxyState>,
|
||||
addr: RwLock<Option<SocketAddr>>,
|
||||
shutdown_tx: RwLock<Option<tokio::sync::oneshot::Sender<()>>>,
|
||||
}
|
||||
|
||||
impl HttpProxy {
|
||||
/// Create a new HTTP proxy.
|
||||
pub fn new(
|
||||
decider: Arc<dyn NetworkPolicyDecider>,
|
||||
credential_resolver: Arc<dyn CredentialResolver>,
|
||||
) -> Self {
|
||||
Self {
|
||||
state: Arc::new(ProxyState {
|
||||
decider,
|
||||
credential_resolver,
|
||||
request_count: std::sync::atomic::AtomicU64::new(0),
|
||||
running: std::sync::atomic::AtomicBool::new(false),
|
||||
}),
|
||||
addr: RwLock::new(None),
|
||||
shutdown_tx: RwLock::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Start the proxy server on the given port (0 for auto-assign).
|
||||
pub async fn start(&self, port: u16) -> Result<SocketAddr> {
|
||||
let listener = TcpListener::bind(format!("127.0.0.1:{}", port))
|
||||
.await
|
||||
.map_err(|e| SandboxError::ProxyError {
|
||||
reason: format!("failed to bind: {}", e),
|
||||
})?;
|
||||
|
||||
let addr = listener
|
||||
.local_addr()
|
||||
.map_err(|e| SandboxError::ProxyError {
|
||||
reason: format!("failed to get local addr: {}", e),
|
||||
})?;
|
||||
|
||||
*self.addr.write().await = Some(addr);
|
||||
|
||||
let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel();
|
||||
*self.shutdown_tx.write().await = Some(shutdown_tx);
|
||||
|
||||
self.state
|
||||
.running
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
|
||||
let state = self.state.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
tracing::info!("Sandbox proxy started on {}", addr);
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
accept_result = listener.accept() => {
|
||||
match accept_result {
|
||||
Ok((stream, _)) => {
|
||||
let io = TokioIo::new(stream);
|
||||
let state = state.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let service = service_fn(move |req| {
|
||||
let state = state.clone();
|
||||
async move { handle_request(req, state).await }
|
||||
});
|
||||
|
||||
if let Err(e) = http1::Builder::new()
|
||||
.preserve_header_case(true)
|
||||
.title_case_headers(true)
|
||||
.serve_connection(io, service)
|
||||
.with_upgrades()
|
||||
.await
|
||||
{
|
||||
tracing::debug!("Proxy connection error: {}", e);
|
||||
}
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("Proxy accept error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = &mut shutdown_rx => {
|
||||
tracing::info!("Sandbox proxy shutting down");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
state
|
||||
.running
|
||||
.store(false, std::sync::atomic::Ordering::SeqCst);
|
||||
});
|
||||
|
||||
Ok(addr)
|
||||
}
|
||||
|
||||
/// Stop the proxy server.
|
||||
pub async fn stop(&self) {
|
||||
if let Some(tx) = self.shutdown_tx.write().await.take() {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the address the proxy is listening on.
|
||||
pub async fn addr(&self) -> Option<SocketAddr> {
|
||||
*self.addr.read().await
|
||||
}
|
||||
|
||||
/// Check if the proxy is running.
|
||||
pub fn is_running(&self) -> bool {
|
||||
self.state.running.load(std::sync::atomic::Ordering::SeqCst)
|
||||
}
|
||||
|
||||
/// Get the number of requests handled.
|
||||
pub fn request_count(&self) -> u64 {
|
||||
self.state
|
||||
.request_count
|
||||
.load(std::sync::atomic::Ordering::SeqCst)
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle an incoming proxy request.
|
||||
async fn handle_request(
|
||||
req: Request<hyper::body::Incoming>,
|
||||
state: Arc<ProxyState>,
|
||||
) -> std::result::Result<Response<BoxBody<Bytes, Infallible>>, Infallible> {
|
||||
state
|
||||
.request_count
|
||||
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
|
||||
// Handle CONNECT method for HTTPS tunneling
|
||||
if req.method() == Method::CONNECT {
|
||||
return Ok(handle_connect(req, state).await);
|
||||
}
|
||||
|
||||
// For HTTP requests, validate and forward
|
||||
let uri = req.uri().to_string();
|
||||
let method = req.method().to_string();
|
||||
|
||||
let network_req = match NetworkRequest::from_url(&method, &uri) {
|
||||
Some(r) => r,
|
||||
None => {
|
||||
tracing::warn!("Proxy: invalid URL: {}", uri);
|
||||
return Ok(error_response(
|
||||
StatusCode::BAD_REQUEST,
|
||||
"Invalid URL".to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// Make policy decision
|
||||
let decision = state.decider.decide(&network_req).await;
|
||||
|
||||
match decision {
|
||||
NetworkDecision::Deny { reason } => {
|
||||
tracing::info!("Proxy: blocked {} {} - {}", method, uri, reason);
|
||||
Ok(error_response(StatusCode::FORBIDDEN, reason))
|
||||
}
|
||||
NetworkDecision::Allow | NetworkDecision::AllowWithCredentials { .. } => {
|
||||
// Forward the request
|
||||
forward_request(req, decision, state).await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle CONNECT method for HTTPS tunneling.
|
||||
async fn handle_connect(
|
||||
req: Request<hyper::body::Incoming>,
|
||||
state: Arc<ProxyState>,
|
||||
) -> Response<BoxBody<Bytes, Infallible>> {
|
||||
// Extract host from CONNECT target
|
||||
let host = req.uri().authority().map(|a| a.host().to_string());
|
||||
|
||||
let host = match host {
|
||||
Some(h) => h,
|
||||
None => {
|
||||
return error_response(StatusCode::BAD_REQUEST, "Missing host".to_string());
|
||||
}
|
||||
};
|
||||
|
||||
// Check if host is allowed
|
||||
let network_req = NetworkRequest {
|
||||
method: "CONNECT".to_string(),
|
||||
url: format!("https://{}", host),
|
||||
host: host.clone(),
|
||||
path: "/".to_string(),
|
||||
};
|
||||
|
||||
let decision = state.decider.decide(&network_req).await;
|
||||
|
||||
if !decision.is_allowed() {
|
||||
if let NetworkDecision::Deny { reason } = decision {
|
||||
tracing::info!("Proxy: blocked CONNECT {} - {}", host, reason);
|
||||
return error_response(StatusCode::FORBIDDEN, reason);
|
||||
}
|
||||
}
|
||||
|
||||
tracing::debug!("Proxy: allowing CONNECT to {}", host);
|
||||
|
||||
// For CONNECT, we return 200 OK and the client will upgrade to TLS
|
||||
// The actual TLS connection goes directly to the target, we just act as a tunnel
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(empty_body())
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
/// Forward a request to the target server.
|
||||
async fn forward_request(
|
||||
req: Request<hyper::body::Incoming>,
|
||||
decision: NetworkDecision,
|
||||
state: Arc<ProxyState>,
|
||||
) -> std::result::Result<Response<BoxBody<Bytes, Infallible>>, Infallible> {
|
||||
let method = req.method().clone();
|
||||
let uri = req.uri().clone();
|
||||
|
||||
// Build the forwarded request
|
||||
let client = reqwest::Client::new();
|
||||
let mut builder = client.request(
|
||||
reqwest::Method::from_bytes(method.as_str().as_bytes()).unwrap_or(reqwest::Method::GET),
|
||||
uri.to_string(),
|
||||
);
|
||||
|
||||
// Copy headers (except hop-by-hop headers)
|
||||
for (name, value) in req.headers() {
|
||||
if !is_hop_by_hop_header(name.as_str()) {
|
||||
if let Ok(v) = value.to_str() {
|
||||
builder = builder.header(name.as_str(), v);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Inject credentials if needed
|
||||
if let NetworkDecision::AllowWithCredentials {
|
||||
secret_name,
|
||||
location,
|
||||
} = decision
|
||||
{
|
||||
if let Some(credential) = state.credential_resolver.resolve(&secret_name).await {
|
||||
builder = match location {
|
||||
CredentialLocation::AuthorizationBearer => {
|
||||
builder.header("Authorization", format!("Bearer {}", credential))
|
||||
}
|
||||
CredentialLocation::Header(header_name) => builder.header(header_name, credential),
|
||||
CredentialLocation::QueryParam(param_name) => {
|
||||
builder.query(&[(param_name, credential)])
|
||||
}
|
||||
};
|
||||
tracing::debug!("Proxy: injected credential for {}", secret_name);
|
||||
} else {
|
||||
tracing::warn!("Proxy: credential {} not found", secret_name);
|
||||
}
|
||||
}
|
||||
|
||||
// Copy body
|
||||
let body_bytes = match req.collect().await {
|
||||
Ok(collected) => collected.to_bytes(),
|
||||
Err(e) => {
|
||||
tracing::error!("Proxy: failed to read request body: {}", e);
|
||||
return Ok(error_response(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"Failed to read body".to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
if !body_bytes.is_empty() {
|
||||
builder = builder.body(body_bytes.to_vec());
|
||||
}
|
||||
|
||||
// Send the request
|
||||
match builder.send().await {
|
||||
Ok(response) => {
|
||||
let status = response.status();
|
||||
let headers = response.headers().clone();
|
||||
|
||||
match response.bytes().await {
|
||||
Ok(body) => {
|
||||
let mut builder = Response::builder().status(status.as_u16());
|
||||
|
||||
for (name, value) in headers.iter() {
|
||||
if !is_hop_by_hop_header(name.as_str()) {
|
||||
builder = builder.header(name.as_str(), value.as_bytes());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(builder.body(full_body(body)).unwrap())
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("Proxy: failed to read response body: {}", e);
|
||||
Ok(error_response(
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"Failed to read response".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("Proxy: request failed: {}", e);
|
||||
Ok(error_response(
|
||||
StatusCode::BAD_GATEWAY,
|
||||
format!("Request failed: {}", e),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if a header is hop-by-hop (should not be forwarded).
|
||||
fn is_hop_by_hop_header(name: &str) -> bool {
|
||||
matches!(
|
||||
name.to_lowercase().as_str(),
|
||||
"connection"
|
||||
| "keep-alive"
|
||||
| "proxy-authenticate"
|
||||
| "proxy-authorization"
|
||||
| "te"
|
||||
| "trailers"
|
||||
| "transfer-encoding"
|
||||
| "upgrade"
|
||||
)
|
||||
}
|
||||
|
||||
/// Create an error response.
|
||||
fn error_response(status: StatusCode, message: String) -> Response<BoxBody<Bytes, Infallible>> {
|
||||
Response::builder()
|
||||
.status(status)
|
||||
.header("Content-Type", "text/plain")
|
||||
.body(full_body(Bytes::from(message)))
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
/// Create an empty body.
|
||||
fn empty_body() -> BoxBody<Bytes, Infallible> {
|
||||
Empty::<Bytes>::new().map_err(|_| unreachable!()).boxed()
|
||||
}
|
||||
|
||||
/// Create a body from bytes.
|
||||
fn full_body(bytes: Bytes) -> BoxBody<Bytes, Infallible> {
|
||||
Full::new(bytes).map_err(|_| unreachable!()).boxed()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::sandbox::proxy::allowlist::DomainAllowlist;
|
||||
use crate::sandbox::proxy::policy::DefaultPolicyDecider;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_proxy_starts_and_stops() {
|
||||
let allowlist = DomainAllowlist::new(&["example.com".to_string()]);
|
||||
let decider = Arc::new(DefaultPolicyDecider::new(allowlist, vec![]));
|
||||
let resolver = Arc::new(NoCredentialResolver);
|
||||
|
||||
let proxy = HttpProxy::new(decider, resolver);
|
||||
|
||||
let addr = proxy.start(0).await.unwrap();
|
||||
assert!(proxy.is_running());
|
||||
assert!(addr.port() > 0);
|
||||
|
||||
proxy.stop().await;
|
||||
// Give it a moment to shut down
|
||||
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hop_by_hop_headers() {
|
||||
assert!(is_hop_by_hop_header("connection"));
|
||||
assert!(is_hop_by_hop_header("Connection"));
|
||||
assert!(is_hop_by_hop_header("transfer-encoding"));
|
||||
assert!(!is_hop_by_hop_header("content-type"));
|
||||
assert!(!is_hop_by_hop_header("authorization"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
//! Network proxy for sandboxed container access.
|
||||
//!
|
||||
//! The proxy provides:
|
||||
//! - Domain allowlist validation
|
||||
//! - Credential injection for API calls
|
||||
//! - Request logging and monitoring
|
||||
//!
|
||||
//! # Architecture
|
||||
//!
|
||||
//! ```text
|
||||
//! ┌─────────────────────────────────────────────────────────────────┐
|
||||
//! │ Network Proxy │
|
||||
//! │ │
|
||||
//! │ ┌─────────────┐ ┌─────────────┐ ┌─────────────────────┐ │
|
||||
//! │ │ HTTP Proxy │───▶│ Policy │───▶│ Credential Resolver │ │
|
||||
//! │ │ Server │ │ Decider │ │ │ │
|
||||
//! │ └─────────────┘ └─────────────┘ └─────────────────────┘ │
|
||||
//! │ │ │ │
|
||||
//! │ │ ▼ │
|
||||
//! │ │ ┌─────────────┐ │
|
||||
//! │ │ │ Allowlist │ │
|
||||
//! │ │ │ Validator │ │
|
||||
//! │ │ └─────────────┘ │
|
||||
//! │ ▼ │
|
||||
//! │ ┌──────────────────────────────────────────────────────────┐ │
|
||||
//! │ │ Internet │ │
|
||||
//! │ └──────────────────────────────────────────────────────────┘ │
|
||||
//! └─────────────────────────────────────────────────────────────────┘
|
||||
//! ```
|
||||
|
||||
pub mod allowlist;
|
||||
pub mod http;
|
||||
pub mod policy;
|
||||
|
||||
pub use allowlist::{DomainAllowlist, DomainPattern, DomainValidationResult};
|
||||
pub use http::{CredentialResolver, EnvCredentialResolver, HttpProxy, NoCredentialResolver};
|
||||
pub use policy::{
|
||||
AllowAllDecider, DefaultPolicyDecider, DenyAllDecider, NetworkDecision, NetworkPolicyDecider,
|
||||
NetworkRequest,
|
||||
};
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::sandbox::config::{
|
||||
CredentialMapping, SandboxConfig, SandboxPolicy, default_credential_mappings,
|
||||
};
|
||||
use crate::sandbox::error::Result;
|
||||
|
||||
/// Creates a configured network proxy from sandbox config.
|
||||
pub struct NetworkProxyBuilder {
|
||||
allowlist: Vec<String>,
|
||||
credential_mappings: Vec<CredentialMapping>,
|
||||
credential_resolver: Arc<dyn CredentialResolver>,
|
||||
policy: SandboxPolicy,
|
||||
}
|
||||
|
||||
impl NetworkProxyBuilder {
|
||||
/// Create a new builder with default settings.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
allowlist: crate::sandbox::config::default_allowlist(),
|
||||
credential_mappings: default_credential_mappings(),
|
||||
credential_resolver: Arc::new(EnvCredentialResolver),
|
||||
policy: SandboxPolicy::ReadOnly,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create from a sandbox config.
|
||||
pub fn from_config(config: &SandboxConfig) -> Self {
|
||||
Self {
|
||||
allowlist: config.network_allowlist.clone(),
|
||||
credential_mappings: default_credential_mappings(),
|
||||
credential_resolver: Arc::new(EnvCredentialResolver),
|
||||
policy: config.policy,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the domain allowlist.
|
||||
pub fn with_allowlist(mut self, domains: Vec<String>) -> Self {
|
||||
self.allowlist = domains;
|
||||
self
|
||||
}
|
||||
|
||||
/// Add a domain to the allowlist.
|
||||
pub fn allow_domain(mut self, domain: &str) -> Self {
|
||||
self.allowlist.push(domain.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
/// Set credential mappings.
|
||||
pub fn with_credentials(mut self, mappings: Vec<CredentialMapping>) -> Self {
|
||||
self.credential_mappings = mappings;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the credential resolver.
|
||||
pub fn with_credential_resolver(mut self, resolver: Arc<dyn CredentialResolver>) -> Self {
|
||||
self.credential_resolver = resolver;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the sandbox policy.
|
||||
pub fn with_policy(mut self, policy: SandboxPolicy) -> Self {
|
||||
self.policy = policy;
|
||||
self
|
||||
}
|
||||
|
||||
/// Build the HTTP proxy.
|
||||
pub fn build(self) -> HttpProxy {
|
||||
let decider: Arc<dyn NetworkPolicyDecider> = if self.policy.has_full_network() {
|
||||
Arc::new(AllowAllDecider)
|
||||
} else {
|
||||
Arc::new(DefaultPolicyDecider::new(
|
||||
DomainAllowlist::new(&self.allowlist),
|
||||
self.credential_mappings,
|
||||
))
|
||||
};
|
||||
|
||||
HttpProxy::new(decider, self.credential_resolver)
|
||||
}
|
||||
|
||||
/// Build and start the proxy on the given port.
|
||||
pub async fn build_and_start(self, port: u16) -> Result<HttpProxy> {
|
||||
let proxy = self.build();
|
||||
proxy.start(port).await?;
|
||||
Ok(proxy)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for NetworkProxyBuilder {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_builder_default() {
|
||||
let builder = NetworkProxyBuilder::new();
|
||||
assert!(!builder.allowlist.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_builder_with_custom_allowlist() {
|
||||
let builder = NetworkProxyBuilder::new()
|
||||
.with_allowlist(vec!["custom.com".to_string()])
|
||||
.allow_domain("another.com");
|
||||
|
||||
assert!(builder.allowlist.contains(&"custom.com".to_string()));
|
||||
assert!(builder.allowlist.contains(&"another.com".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_builder_builds_proxy() {
|
||||
let proxy = NetworkProxyBuilder::new()
|
||||
.with_policy(SandboxPolicy::ReadOnly)
|
||||
.build();
|
||||
|
||||
assert!(!proxy.is_running());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,229 @@
|
||||
//! Network policy decision making.
|
||||
//!
|
||||
//! Determines whether network requests should be allowed, denied,
|
||||
//! or allowed with credential injection.
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::sandbox::config::{CredentialLocation, CredentialMapping};
|
||||
use crate::sandbox::proxy::allowlist::DomainAllowlist;
|
||||
|
||||
/// A network request to be evaluated.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct NetworkRequest {
|
||||
/// HTTP method (GET, POST, etc.).
|
||||
pub method: String,
|
||||
/// Full URL being requested.
|
||||
pub url: String,
|
||||
/// Host extracted from URL.
|
||||
pub host: String,
|
||||
/// Path portion of the URL.
|
||||
pub path: String,
|
||||
}
|
||||
|
||||
impl NetworkRequest {
|
||||
/// Create from a URL string.
|
||||
pub fn from_url(method: &str, url: &str) -> Option<Self> {
|
||||
let host = crate::sandbox::proxy::allowlist::extract_host(url)?;
|
||||
let path = extract_path(url);
|
||||
|
||||
Some(Self {
|
||||
method: method.to_uppercase(),
|
||||
url: url.to_string(),
|
||||
host,
|
||||
path,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract path from a URL.
|
||||
fn extract_path(url: &str) -> String {
|
||||
// Find the start of the path (after ://)
|
||||
if let Some(idx) = url.find("://") {
|
||||
let rest = &url[idx + 3..];
|
||||
if let Some(path_start) = rest.find('/') {
|
||||
return rest[path_start..].to_string();
|
||||
}
|
||||
}
|
||||
"/".to_string()
|
||||
}
|
||||
|
||||
/// Decision for a network request.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum NetworkDecision {
|
||||
/// Allow the request as-is.
|
||||
Allow,
|
||||
/// Allow with credential injection.
|
||||
AllowWithCredentials {
|
||||
/// Name of the secret to look up.
|
||||
secret_name: String,
|
||||
/// Where to inject the credential.
|
||||
location: CredentialLocation,
|
||||
},
|
||||
/// Deny the request.
|
||||
Deny {
|
||||
/// Reason for denial.
|
||||
reason: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl NetworkDecision {
|
||||
pub fn is_allowed(&self) -> bool {
|
||||
!matches!(self, NetworkDecision::Deny { .. })
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for making network policy decisions.
|
||||
#[async_trait]
|
||||
pub trait NetworkPolicyDecider: Send + Sync {
|
||||
/// Decide whether a request should be allowed.
|
||||
async fn decide(&self, request: &NetworkRequest) -> NetworkDecision;
|
||||
}
|
||||
|
||||
/// Default policy decider that uses allowlist and credential mappings.
|
||||
pub struct DefaultPolicyDecider {
|
||||
allowlist: DomainAllowlist,
|
||||
credential_mappings: Vec<CredentialMapping>,
|
||||
}
|
||||
|
||||
impl DefaultPolicyDecider {
|
||||
/// Create a new policy decider.
|
||||
pub fn new(allowlist: DomainAllowlist, credential_mappings: Vec<CredentialMapping>) -> Self {
|
||||
Self {
|
||||
allowlist,
|
||||
credential_mappings,
|
||||
}
|
||||
}
|
||||
|
||||
/// Find credential mapping for a domain.
|
||||
fn find_credential(&self, host: &str) -> Option<&CredentialMapping> {
|
||||
let host_lower = host.to_lowercase();
|
||||
self.credential_mappings
|
||||
.iter()
|
||||
.find(|m| m.domain.to_lowercase() == host_lower)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl NetworkPolicyDecider for DefaultPolicyDecider {
|
||||
async fn decide(&self, request: &NetworkRequest) -> NetworkDecision {
|
||||
// First check if the domain is allowed
|
||||
let validation = self.allowlist.is_allowed(&request.host);
|
||||
if !validation.is_allowed() {
|
||||
if let crate::sandbox::proxy::allowlist::DomainValidationResult::Denied(reason) =
|
||||
validation
|
||||
{
|
||||
return NetworkDecision::Deny { reason };
|
||||
}
|
||||
}
|
||||
|
||||
// Check if we need to inject credentials
|
||||
if let Some(mapping) = self.find_credential(&request.host) {
|
||||
return NetworkDecision::AllowWithCredentials {
|
||||
secret_name: mapping.secret_name.clone(),
|
||||
location: mapping.location.clone(),
|
||||
};
|
||||
}
|
||||
|
||||
NetworkDecision::Allow
|
||||
}
|
||||
}
|
||||
|
||||
/// A policy decider that allows everything (use with FullAccess policy).
|
||||
pub struct AllowAllDecider;
|
||||
|
||||
#[async_trait]
|
||||
impl NetworkPolicyDecider for AllowAllDecider {
|
||||
async fn decide(&self, _request: &NetworkRequest) -> NetworkDecision {
|
||||
NetworkDecision::Allow
|
||||
}
|
||||
}
|
||||
|
||||
/// A policy decider that denies everything.
|
||||
pub struct DenyAllDecider {
|
||||
reason: String,
|
||||
}
|
||||
|
||||
impl DenyAllDecider {
|
||||
pub fn new(reason: &str) -> Self {
|
||||
Self {
|
||||
reason: reason.to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl NetworkPolicyDecider for DenyAllDecider {
|
||||
async fn decide(&self, _request: &NetworkRequest) -> NetworkDecision {
|
||||
NetworkDecision::Deny {
|
||||
reason: self.reason.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_network_request_from_url() {
|
||||
let req = NetworkRequest::from_url("GET", "https://api.example.com/v1/data").unwrap();
|
||||
assert_eq!(req.method, "GET");
|
||||
assert_eq!(req.host, "api.example.com");
|
||||
assert_eq!(req.path, "/v1/data");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_path() {
|
||||
assert_eq!(
|
||||
extract_path("https://example.com/api/v1"),
|
||||
"/api/v1".to_string()
|
||||
);
|
||||
assert_eq!(extract_path("https://example.com"), "/".to_string());
|
||||
assert_eq!(extract_path("https://example.com/"), "/".to_string());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_default_policy_allows_listed_domain() {
|
||||
let allowlist = DomainAllowlist::new(&["crates.io".to_string()]);
|
||||
let decider = DefaultPolicyDecider::new(allowlist, vec![]);
|
||||
|
||||
let req = NetworkRequest::from_url("GET", "https://crates.io/api/v1/crates").unwrap();
|
||||
let decision = decider.decide(&req).await;
|
||||
|
||||
assert!(decision.is_allowed());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_default_policy_denies_unlisted_domain() {
|
||||
let allowlist = DomainAllowlist::new(&["crates.io".to_string()]);
|
||||
let decider = DefaultPolicyDecider::new(allowlist, vec![]);
|
||||
|
||||
let req = NetworkRequest::from_url("GET", "https://evil.com/steal").unwrap();
|
||||
let decision = decider.decide(&req).await;
|
||||
|
||||
assert!(!decision.is_allowed());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_credential_injection() {
|
||||
let allowlist = DomainAllowlist::new(&["api.openai.com".to_string()]);
|
||||
let credentials = vec![CredentialMapping {
|
||||
domain: "api.openai.com".to_string(),
|
||||
secret_name: "OPENAI_API_KEY".to_string(),
|
||||
location: CredentialLocation::AuthorizationBearer,
|
||||
}];
|
||||
let decider = DefaultPolicyDecider::new(allowlist, credentials);
|
||||
|
||||
let req =
|
||||
NetworkRequest::from_url("POST", "https://api.openai.com/v1/chat/completions").unwrap();
|
||||
let decision = decider.decide(&req).await;
|
||||
|
||||
match decision {
|
||||
NetworkDecision::AllowWithCredentials { secret_name, .. } => {
|
||||
assert_eq!(secret_name, "OPENAI_API_KEY");
|
||||
}
|
||||
_ => panic!("Expected AllowWithCredentials"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,417 @@
|
||||
//! Job management tools.
|
||||
//!
|
||||
//! These tools allow the LLM to manage jobs:
|
||||
//! - Create new jobs/tasks
|
||||
//! - List existing jobs
|
||||
//! - Check job status
|
||||
//! - Cancel running jobs
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::context::{ContextManager, JobContext, JobState};
|
||||
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||
|
||||
/// Tool for creating a new job.
|
||||
pub struct CreateJobTool {
|
||||
context_manager: Arc<ContextManager>,
|
||||
}
|
||||
|
||||
impl CreateJobTool {
|
||||
pub fn new(context_manager: Arc<ContextManager>) -> Self {
|
||||
Self { context_manager }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for CreateJobTool {
|
||||
fn name(&self) -> &str {
|
||||
"create_job"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Create a new job or task for the agent to work on. Use this when the user wants \
|
||||
you to do something substantial that should be tracked as a separate job."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"title": {
|
||||
"type": "string",
|
||||
"description": "A short title for the job (max 100 chars)"
|
||||
},
|
||||
"description": {
|
||||
"type": "string",
|
||||
"description": "Full description of what needs to be done"
|
||||
}
|
||||
},
|
||||
"required": ["title", "description"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
let title = params
|
||||
.get("title")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidParameters("missing 'title' parameter".into()))?;
|
||||
|
||||
let description = params
|
||||
.get("description")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| {
|
||||
ToolError::InvalidParameters("missing 'description' parameter".into())
|
||||
})?;
|
||||
|
||||
match self.context_manager.create_job(title, description).await {
|
||||
Ok(job_id) => {
|
||||
let result = serde_json::json!({
|
||||
"job_id": job_id.to_string(),
|
||||
"title": title,
|
||||
"status": "pending",
|
||||
"message": format!("Created job '{}'", title)
|
||||
});
|
||||
Ok(ToolOutput::success(result, start.elapsed()))
|
||||
}
|
||||
Err(e) => {
|
||||
let result = serde_json::json!({
|
||||
"error": e.to_string()
|
||||
});
|
||||
Ok(ToolOutput::success(result, start.elapsed()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Tool for listing jobs.
|
||||
pub struct ListJobsTool {
|
||||
context_manager: Arc<ContextManager>,
|
||||
}
|
||||
|
||||
impl ListJobsTool {
|
||||
pub fn new(context_manager: Arc<ContextManager>) -> Self {
|
||||
Self { context_manager }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ListJobsTool {
|
||||
fn name(&self) -> &str {
|
||||
"list_jobs"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"List all jobs or filter by status. Shows job IDs, titles, and current status."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"filter": {
|
||||
"type": "string",
|
||||
"description": "Filter by status: 'active', 'completed', 'failed', 'all' (default: 'all')",
|
||||
"enum": ["active", "completed", "failed", "all"]
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
let filter = params
|
||||
.get("filter")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("all");
|
||||
|
||||
let job_ids = match filter {
|
||||
"active" => self.context_manager.active_jobs().await,
|
||||
_ => self.context_manager.all_jobs().await,
|
||||
};
|
||||
|
||||
let mut jobs = Vec::new();
|
||||
for job_id in job_ids {
|
||||
if let Ok(ctx) = self.context_manager.get_context(job_id).await {
|
||||
let include = match filter {
|
||||
"completed" => ctx.state == JobState::Completed,
|
||||
"failed" => ctx.state == JobState::Failed,
|
||||
"active" => ctx.state.is_active(),
|
||||
_ => true,
|
||||
};
|
||||
|
||||
if include {
|
||||
jobs.push(serde_json::json!({
|
||||
"job_id": job_id.to_string(),
|
||||
"title": ctx.title,
|
||||
"status": format!("{:?}", ctx.state),
|
||||
"created_at": ctx.created_at.to_rfc3339()
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let summary = self.context_manager.summary().await;
|
||||
|
||||
let result = serde_json::json!({
|
||||
"jobs": jobs,
|
||||
"summary": {
|
||||
"total": summary.total,
|
||||
"pending": summary.pending,
|
||||
"in_progress": summary.in_progress,
|
||||
"completed": summary.completed,
|
||||
"failed": summary.failed
|
||||
}
|
||||
});
|
||||
|
||||
Ok(ToolOutput::success(result, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Tool for checking job status.
|
||||
pub struct JobStatusTool {
|
||||
context_manager: Arc<ContextManager>,
|
||||
}
|
||||
|
||||
impl JobStatusTool {
|
||||
pub fn new(context_manager: Arc<ContextManager>) -> Self {
|
||||
Self { context_manager }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for JobStatusTool {
|
||||
fn name(&self) -> &str {
|
||||
"job_status"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Check the status and details of a specific job by its ID."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"job_id": {
|
||||
"type": "string",
|
||||
"description": "The UUID of the job to check"
|
||||
}
|
||||
},
|
||||
"required": ["job_id"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
let job_id_str = params
|
||||
.get("job_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidParameters("missing 'job_id' parameter".into()))?;
|
||||
|
||||
let job_id = Uuid::parse_str(job_id_str).map_err(|_| {
|
||||
ToolError::InvalidParameters(format!("invalid job ID format: {}", job_id_str))
|
||||
})?;
|
||||
|
||||
match self.context_manager.get_context(job_id).await {
|
||||
Ok(ctx) => {
|
||||
let result = serde_json::json!({
|
||||
"job_id": job_id.to_string(),
|
||||
"title": ctx.title,
|
||||
"description": ctx.description,
|
||||
"status": format!("{:?}", ctx.state),
|
||||
"created_at": ctx.created_at.to_rfc3339(),
|
||||
"started_at": ctx.started_at.map(|t| t.to_rfc3339()),
|
||||
"completed_at": ctx.completed_at.map(|t| t.to_rfc3339()),
|
||||
"actual_cost": ctx.actual_cost.to_string()
|
||||
});
|
||||
Ok(ToolOutput::success(result, start.elapsed()))
|
||||
}
|
||||
Err(e) => {
|
||||
let result = serde_json::json!({
|
||||
"error": format!("Job not found: {}", e)
|
||||
});
|
||||
Ok(ToolOutput::success(result, start.elapsed()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Tool for canceling a job.
|
||||
pub struct CancelJobTool {
|
||||
context_manager: Arc<ContextManager>,
|
||||
}
|
||||
|
||||
impl CancelJobTool {
|
||||
pub fn new(context_manager: Arc<ContextManager>) -> Self {
|
||||
Self { context_manager }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for CancelJobTool {
|
||||
fn name(&self) -> &str {
|
||||
"cancel_job"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Cancel a running or pending job. The job will be marked as cancelled and stopped."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"job_id": {
|
||||
"type": "string",
|
||||
"description": "The UUID of the job to cancel"
|
||||
}
|
||||
},
|
||||
"required": ["job_id"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
let job_id_str = params
|
||||
.get("job_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidParameters("missing 'job_id' parameter".into()))?;
|
||||
|
||||
let job_id = Uuid::parse_str(job_id_str).map_err(|_| {
|
||||
ToolError::InvalidParameters(format!("invalid job ID format: {}", job_id_str))
|
||||
})?;
|
||||
|
||||
// Transition to cancelled state
|
||||
match self
|
||||
.context_manager
|
||||
.update_context(job_id, |ctx| {
|
||||
ctx.transition_to(JobState::Cancelled, Some("Cancelled by user".to_string()))
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(Ok(())) => {
|
||||
let result = serde_json::json!({
|
||||
"job_id": job_id.to_string(),
|
||||
"status": "cancelled",
|
||||
"message": "Job cancelled successfully"
|
||||
});
|
||||
Ok(ToolOutput::success(result, start.elapsed()))
|
||||
}
|
||||
Ok(Err(reason)) => {
|
||||
let result = serde_json::json!({
|
||||
"error": format!("Cannot cancel job: {}", reason)
|
||||
});
|
||||
Ok(ToolOutput::success(result, start.elapsed()))
|
||||
}
|
||||
Err(e) => {
|
||||
let result = serde_json::json!({
|
||||
"error": format!("Job not found: {}", e)
|
||||
});
|
||||
Ok(ToolOutput::success(result, start.elapsed()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
true // Canceling a job should require approval
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_job_tool() {
|
||||
let manager = Arc::new(ContextManager::new(5));
|
||||
let tool = CreateJobTool::new(manager.clone());
|
||||
|
||||
let params = serde_json::json!({
|
||||
"title": "Test Job",
|
||||
"description": "A test job description"
|
||||
});
|
||||
|
||||
let ctx = JobContext::default();
|
||||
let result = tool.execute(params, &ctx).await.unwrap();
|
||||
|
||||
let job_id = result.result.get("job_id").unwrap().as_str().unwrap();
|
||||
assert!(!job_id.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_list_jobs_tool() {
|
||||
let manager = Arc::new(ContextManager::new(5));
|
||||
|
||||
// Create some jobs
|
||||
manager.create_job("Job 1", "Desc 1").await.unwrap();
|
||||
manager.create_job("Job 2", "Desc 2").await.unwrap();
|
||||
|
||||
let tool = ListJobsTool::new(manager);
|
||||
|
||||
let params = serde_json::json!({});
|
||||
let ctx = JobContext::default();
|
||||
let result = tool.execute(params, &ctx).await.unwrap();
|
||||
|
||||
let jobs = result.result.get("jobs").unwrap().as_array().unwrap();
|
||||
assert_eq!(jobs.len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_job_status_tool() {
|
||||
let manager = Arc::new(ContextManager::new(5));
|
||||
let job_id = manager.create_job("Test Job", "Description").await.unwrap();
|
||||
|
||||
let tool = JobStatusTool::new(manager);
|
||||
|
||||
let params = serde_json::json!({
|
||||
"job_id": job_id.to_string()
|
||||
});
|
||||
let ctx = JobContext::default();
|
||||
let result = tool.execute(params, &ctx).await.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
result.result.get("title").unwrap().as_str().unwrap(),
|
||||
"Test Job"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -4,6 +4,7 @@ mod echo;
|
||||
mod ecommerce;
|
||||
mod file;
|
||||
mod http;
|
||||
mod job;
|
||||
mod json;
|
||||
mod marketplace;
|
||||
mod memory;
|
||||
@@ -16,6 +17,7 @@ pub use echo::EchoTool;
|
||||
pub use ecommerce::EcommerceTool;
|
||||
pub use file::{ApplyPatchTool, ListDirTool, ReadFileTool, WriteFileTool};
|
||||
pub use http::HttpTool;
|
||||
pub use job::{CancelJobTool, CreateJobTool, JobStatusTool, ListJobsTool};
|
||||
pub use json::JsonTool;
|
||||
pub use marketplace::MarketplaceTool;
|
||||
pub use memory::{MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool};
|
||||
|
||||
+145
-30
@@ -1,15 +1,27 @@
|
||||
//! Shell execution tool for running commands in a sandboxed environment.
|
||||
//!
|
||||
//! Provides controlled command execution with:
|
||||
//! - Docker sandbox isolation (when enabled)
|
||||
//! - Working directory isolation
|
||||
//! - Timeout enforcement
|
||||
//! - Output capture and truncation
|
||||
//! - Blocked command patterns for safety
|
||||
//!
|
||||
//! # Execution Modes
|
||||
//!
|
||||
//! When sandbox is available and enabled:
|
||||
//! - Commands run inside ephemeral Docker containers
|
||||
//! - Network traffic goes through a validating proxy
|
||||
//! - Credentials are injected by the proxy, never exposed to commands
|
||||
//!
|
||||
//! When sandbox is unavailable:
|
||||
//! - Commands run directly on host with basic protections
|
||||
//! - Blocked command patterns are still enforced
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::path::PathBuf;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::Stdio;
|
||||
use std::sync::LazyLock;
|
||||
use std::sync::{Arc, LazyLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
@@ -17,6 +29,7 @@ use tokio::io::AsyncReadExt;
|
||||
use tokio::process::Command;
|
||||
|
||||
use crate::context::JobContext;
|
||||
use crate::sandbox::{SandboxManager, SandboxPolicy};
|
||||
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||
|
||||
/// Maximum output size before truncation (64KB).
|
||||
@@ -62,7 +75,6 @@ static DANGEROUS_PATTERNS: LazyLock<Vec<&'static str>> = LazyLock::new(|| {
|
||||
});
|
||||
|
||||
/// Shell command execution tool.
|
||||
#[derive(Debug)]
|
||||
pub struct ShellTool {
|
||||
/// Working directory for commands (if None, uses job's working dir or cwd).
|
||||
working_dir: Option<PathBuf>,
|
||||
@@ -70,6 +82,22 @@ pub struct ShellTool {
|
||||
timeout: Duration,
|
||||
/// Whether to allow potentially dangerous commands (requires explicit approval).
|
||||
allow_dangerous: bool,
|
||||
/// Optional sandbox manager for Docker execution.
|
||||
sandbox: Option<Arc<SandboxManager>>,
|
||||
/// Sandbox policy to use when sandbox is available.
|
||||
sandbox_policy: SandboxPolicy,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for ShellTool {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("ShellTool")
|
||||
.field("working_dir", &self.working_dir)
|
||||
.field("timeout", &self.timeout)
|
||||
.field("allow_dangerous", &self.allow_dangerous)
|
||||
.field("sandbox", &self.sandbox.is_some())
|
||||
.field("sandbox_policy", &self.sandbox_policy)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl ShellTool {
|
||||
@@ -79,6 +107,8 @@ impl ShellTool {
|
||||
working_dir: None,
|
||||
timeout: DEFAULT_TIMEOUT,
|
||||
allow_dangerous: false,
|
||||
sandbox: None,
|
||||
sandbox_policy: SandboxPolicy::ReadOnly,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -94,6 +124,18 @@ impl ShellTool {
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable sandbox execution with the given manager.
|
||||
pub fn with_sandbox(mut self, sandbox: Arc<SandboxManager>) -> Self {
|
||||
self.sandbox = Some(sandbox);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the sandbox policy.
|
||||
pub fn with_sandbox_policy(mut self, policy: SandboxPolicy) -> Self {
|
||||
self.sandbox_policy = policy;
|
||||
self
|
||||
}
|
||||
|
||||
/// Check if a command is blocked.
|
||||
fn is_blocked(&self, cmd: &str) -> Option<&'static str> {
|
||||
let normalized = cmd.to_lowercase();
|
||||
@@ -115,28 +157,44 @@ impl ShellTool {
|
||||
None
|
||||
}
|
||||
|
||||
/// Execute a command and capture output.
|
||||
async fn execute_command(
|
||||
/// Execute a command through the sandbox.
|
||||
async fn execute_sandboxed(
|
||||
&self,
|
||||
sandbox: &SandboxManager,
|
||||
cmd: &str,
|
||||
workdir: &Path,
|
||||
timeout: Duration,
|
||||
) -> Result<(String, i64), ToolError> {
|
||||
// Override sandbox config timeout if needed
|
||||
let result = tokio::time::timeout(timeout, async {
|
||||
sandbox
|
||||
.execute_with_policy(
|
||||
cmd,
|
||||
workdir,
|
||||
self.sandbox_policy,
|
||||
std::collections::HashMap::new(),
|
||||
)
|
||||
.await
|
||||
})
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(output)) => {
|
||||
let combined = truncate_output(&output.output);
|
||||
Ok((combined, output.exit_code))
|
||||
}
|
||||
Ok(Err(e)) => Err(ToolError::ExecutionFailed(format!("Sandbox error: {}", e))),
|
||||
Err(_) => Err(ToolError::Timeout(timeout)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Execute a command directly (fallback when sandbox unavailable).
|
||||
async fn execute_direct(
|
||||
&self,
|
||||
cmd: &str,
|
||||
workdir: Option<&str>,
|
||||
timeout: Option<u64>,
|
||||
workdir: &PathBuf,
|
||||
timeout: Duration,
|
||||
) -> Result<(String, i32), ToolError> {
|
||||
// Check for blocked commands
|
||||
if let Some(reason) = self.is_blocked(cmd) {
|
||||
return Err(ToolError::NotAuthorized(format!(
|
||||
"{}: {}",
|
||||
reason,
|
||||
truncate_for_error(cmd)
|
||||
)));
|
||||
}
|
||||
|
||||
// Determine working directory
|
||||
let cwd = workdir
|
||||
.map(PathBuf::from)
|
||||
.or_else(|| self.working_dir.clone())
|
||||
.unwrap_or_else(|| std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")));
|
||||
|
||||
// Build command
|
||||
let mut command = if cfg!(target_os = "windows") {
|
||||
let mut c = Command::new("cmd");
|
||||
@@ -149,7 +207,7 @@ impl ShellTool {
|
||||
};
|
||||
|
||||
command
|
||||
.current_dir(&cwd)
|
||||
.current_dir(workdir)
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped());
|
||||
@@ -159,11 +217,8 @@ impl ShellTool {
|
||||
.spawn()
|
||||
.map_err(|e| ToolError::ExecutionFailed(format!("Failed to spawn command: {}", e)))?;
|
||||
|
||||
// Determine timeout
|
||||
let timeout_duration = timeout.map(Duration::from_secs).unwrap_or(self.timeout);
|
||||
|
||||
// Wait with timeout
|
||||
let result = tokio::time::timeout(timeout_duration, async {
|
||||
let result = tokio::time::timeout(timeout, async {
|
||||
let status = child.wait().await?;
|
||||
|
||||
// Read stdout
|
||||
@@ -204,10 +259,56 @@ impl ShellTool {
|
||||
Err(_) => {
|
||||
// Timeout - try to kill the process
|
||||
let _ = child.kill().await;
|
||||
Err(ToolError::Timeout(timeout_duration))
|
||||
Err(ToolError::Timeout(timeout))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Execute a command, using sandbox if available.
|
||||
async fn execute_command(
|
||||
&self,
|
||||
cmd: &str,
|
||||
workdir: Option<&str>,
|
||||
timeout: Option<u64>,
|
||||
) -> Result<(String, i64), ToolError> {
|
||||
// Check for blocked commands
|
||||
if let Some(reason) = self.is_blocked(cmd) {
|
||||
return Err(ToolError::NotAuthorized(format!(
|
||||
"{}: {}",
|
||||
reason,
|
||||
truncate_for_error(cmd)
|
||||
)));
|
||||
}
|
||||
|
||||
// Determine working directory
|
||||
let cwd = workdir
|
||||
.map(PathBuf::from)
|
||||
.or_else(|| self.working_dir.clone())
|
||||
.unwrap_or_else(|| std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")));
|
||||
|
||||
// Determine timeout
|
||||
let timeout_duration = timeout.map(Duration::from_secs).unwrap_or(self.timeout);
|
||||
|
||||
// Try sandbox execution if available
|
||||
if let Some(ref sandbox) = self.sandbox {
|
||||
if sandbox.is_initialized() || sandbox.config().enabled {
|
||||
match self
|
||||
.execute_sandboxed(sandbox, cmd, &cwd, timeout_duration)
|
||||
.await
|
||||
{
|
||||
Ok((output, code)) => return Ok((output, code)),
|
||||
Err(e) => {
|
||||
// Log sandbox failure and fall through to direct execution
|
||||
tracing::warn!("Sandbox execution failed, falling back to direct: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to direct execution
|
||||
let (output, code) = self.execute_direct(cmd, &cwd, timeout_duration).await?;
|
||||
Ok((output, code as i64))
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ShellTool {
|
||||
@@ -224,7 +325,8 @@ impl Tool for ShellTool {
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Execute shell commands. Use for running builds, tests, git operations, and other CLI tasks. \
|
||||
Commands run in a subprocess with captured output. Long-running commands have a timeout."
|
||||
Commands run in a subprocess with captured output. Long-running commands have a timeout. \
|
||||
When Docker sandbox is enabled, commands run in isolated containers for security."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
@@ -265,10 +367,13 @@ impl Tool for ShellTool {
|
||||
let (output, exit_code) = self.execute_command(command, workdir, timeout).await?;
|
||||
let duration = start.elapsed();
|
||||
|
||||
let sandboxed = self.sandbox.is_some();
|
||||
|
||||
let result = serde_json::json!({
|
||||
"output": output,
|
||||
"exit_code": exit_code,
|
||||
"success": exit_code == 0
|
||||
"success": exit_code == 0,
|
||||
"sandboxed": sandboxed
|
||||
});
|
||||
|
||||
Ok(ToolOutput::success(result, duration))
|
||||
@@ -348,4 +453,14 @@ mod tests {
|
||||
|
||||
assert!(matches!(result, Err(ToolError::Timeout(_))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sandbox_policy_builder() {
|
||||
let tool = ShellTool::new()
|
||||
.with_sandbox_policy(SandboxPolicy::WorkspaceWrite)
|
||||
.with_timeout(Duration::from_secs(60));
|
||||
|
||||
assert_eq!(tool.sandbox_policy, SandboxPolicy::WorkspaceWrite);
|
||||
assert_eq!(tool.timeout, Duration::from_secs(60));
|
||||
}
|
||||
}
|
||||
|
||||
+17
-2
@@ -5,12 +5,14 @@ use std::sync::Arc;
|
||||
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::context::ContextManager;
|
||||
use crate::llm::{LlmProvider, ToolDefinition};
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tools::builder::{BuildSoftwareTool, BuilderConfig, LlmSoftwareBuilder};
|
||||
use crate::tools::builtin::{
|
||||
ApplyPatchTool, EchoTool, HttpTool, JsonTool, ListDirTool, MemoryReadTool, MemorySearchTool,
|
||||
MemoryTreeTool, MemoryWriteTool, ReadFileTool, ShellTool, TimeTool, WriteFileTool,
|
||||
ApplyPatchTool, CancelJobTool, CreateJobTool, EchoTool, HttpTool, JobStatusTool, JsonTool,
|
||||
ListDirTool, ListJobsTool, MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool,
|
||||
ReadFileTool, ShellTool, TimeTool, WriteFileTool,
|
||||
};
|
||||
use crate::tools::tool::Tool;
|
||||
use crate::tools::wasm::{
|
||||
@@ -144,6 +146,19 @@ impl ToolRegistry {
|
||||
tracing::info!("Registered 4 memory tools");
|
||||
}
|
||||
|
||||
/// Register job management tools.
|
||||
///
|
||||
/// Job tools allow the LLM to create, list, check status, and cancel jobs.
|
||||
/// These enable natural language job management without hardcoded intent parsing.
|
||||
pub fn register_job_tools(&self, context_manager: Arc<ContextManager>) {
|
||||
self.register_sync(Arc::new(CreateJobTool::new(Arc::clone(&context_manager))));
|
||||
self.register_sync(Arc::new(ListJobsTool::new(Arc::clone(&context_manager))));
|
||||
self.register_sync(Arc::new(JobStatusTool::new(Arc::clone(&context_manager))));
|
||||
self.register_sync(Arc::new(CancelJobTool::new(context_manager)));
|
||||
|
||||
tracing::info!("Registered 4 job management tools");
|
||||
}
|
||||
|
||||
/// Register the software builder tool.
|
||||
///
|
||||
/// The builder tool allows the agent to create new software including WASM tools,
|
||||
|
||||
Reference in New Issue
Block a user