Merge remote-tracking branch 'origin/main'

This commit is contained in:
Illia Polosukhin
2026-02-04 22:11:37 -08:00
24 changed files with 4015 additions and 171 deletions
+15 -14
View File
@@ -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
View File
@@ -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
View File
@@ -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> {
+1
View File
@@ -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;
+8
View File
@@ -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...");
+258
View File
@@ -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()));
}
}
+519
View File
@@ -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;
}
}
+58
View File
@@ -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>;
+474
View File
@@ -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"));
}
}
+113
View File
@@ -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()
}
+251
View File
@@ -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);
}
}
+444
View File
@@ -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"));
}
}
+164
View File
@@ -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());
}
}
+229
View File
@@ -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"),
}
}
}
+417
View File
@@ -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"
);
}
}
+2
View File
@@ -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
View File
@@ -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
View File
@@ -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,