use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use serde_json::{Value, json};
use adk_core::ToolContext;
use crate::backend::SandboxBackend;
use crate::error::SandboxError;
use crate::types::{ExecRequest, Language};
pub struct SandboxTool {
backend: Arc<dyn SandboxBackend>,
}
impl SandboxTool {
pub fn new(backend: Arc<dyn SandboxBackend>) -> Self {
Self { backend }
}
}
const DEFAULT_TIMEOUT_SECS: u64 = 30;
const REQUIRED_SCOPES: &[&str] = &["code:execute"];
fn parse_language(value: &Value) -> Result<Language, String> {
let s = value.as_str().ok_or_else(|| "\"language\" must be a string".to_string())?;
match s {
"rust" => Ok(Language::Rust),
"python" => Ok(Language::Python),
"javascript" => Ok(Language::JavaScript),
"typescript" => Ok(Language::TypeScript),
"wasm" => Ok(Language::Wasm),
"command" => Ok(Language::Command),
other => Err(format!(
"unsupported language \"{other}\". Expected one of: rust, python, javascript, typescript, wasm, command"
)),
}
}
fn sandbox_error_to_json(err: &SandboxError) -> Value {
match err {
SandboxError::Timeout { timeout } => json!({
"status": "timeout",
"stderr": format!("execution timed out after {timeout:?}"),
"duration_ms": timeout.as_millis() as u64,
}),
SandboxError::MemoryExceeded { limit_mb } => json!({
"status": "memory_exceeded",
"stderr": format!("memory limit exceeded: {limit_mb} MB"),
}),
SandboxError::ExecutionFailed(msg) => json!({
"status": "error",
"stderr": msg,
}),
SandboxError::InvalidRequest(msg) => json!({
"status": "error",
"stderr": msg,
}),
SandboxError::BackendUnavailable(msg) => json!({
"status": "error",
"stderr": msg,
}),
SandboxError::EnforcerFailed { enforcer, message } => json!({
"status": "error",
"stderr": format!("enforcer '{enforcer}' failed: {message}"),
}),
SandboxError::EnforcerUnavailable { enforcer, message } => json!({
"status": "error",
"stderr": format!("enforcer '{enforcer}' unavailable: {message}"),
}),
SandboxError::PolicyViolation(msg) => json!({
"status": "error",
"stderr": format!("policy violation: {msg}"),
}),
}
}
#[async_trait]
impl adk_core::Tool for SandboxTool {
fn name(&self) -> &str {
"sandbox_exec"
}
fn description(&self) -> &str {
"Execute code in an isolated sandbox. Supports multiple languages \
including rust, python, javascript, typescript, wasm, and shell commands."
}
fn required_scopes(&self) -> &[&str] {
REQUIRED_SCOPES
}
fn parameters_schema(&self) -> Option<Value> {
Some(json!({
"type": "object",
"properties": {
"language": {
"type": "string",
"enum": ["rust", "python", "javascript", "typescript", "wasm", "command"],
"description": "The programming language of the code to execute."
},
"code": {
"type": "string",
"description": "The source code or command to execute."
},
"stdin": {
"type": "string",
"description": "Optional standard input to feed to the process."
},
"timeout_secs": {
"type": "integer",
"description": "Maximum execution time in seconds.",
"default": DEFAULT_TIMEOUT_SECS,
"minimum": 1,
"maximum": 300
}
},
"required": ["language", "code"]
}))
}
async fn execute(&self, _ctx: Arc<dyn ToolContext>, args: Value) -> adk_core::Result<Value> {
let language = match args.get("language") {
Some(v) => match parse_language(v) {
Ok(lang) => lang,
Err(msg) => {
return Ok(json!({ "status": "error", "stderr": msg }));
}
},
None => {
return Ok(
json!({ "status": "error", "stderr": "missing required field \"language\"" }),
);
}
};
let code = match args.get("code").and_then(|v| v.as_str()) {
Some(c) => c.to_string(),
None => {
return Ok(
json!({ "status": "error", "stderr": "missing required field \"code\"" }),
);
}
};
let stdin = args.get("stdin").and_then(|v| v.as_str()).map(String::from);
let timeout_secs =
args.get("timeout_secs").and_then(|v| v.as_u64()).unwrap_or(DEFAULT_TIMEOUT_SECS);
let request = ExecRequest {
language,
code,
stdin,
timeout: Duration::from_secs(timeout_secs),
memory_limit_mb: None,
env: HashMap::new(),
};
match self.backend.execute(request).await {
Ok(result) => Ok(json!({
"status": "success",
"stdout": result.stdout,
"stderr": result.stderr,
"exit_code": result.exit_code,
"duration_ms": result.duration.as_millis() as u64,
})),
Err(err) => Ok(sandbox_error_to_json(&err)),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::backend::{BackendCapabilities, EnforcedLimits};
use crate::types::ExecResult;
use adk_core::{CallbackContext, Content, EventActions, ReadonlyContext, Tool};
use std::sync::Mutex;
use std::time::Duration;
struct MockBackend {
error: Option<SandboxError>,
result: ExecResult,
}
impl MockBackend {
fn success(stdout: &str, exit_code: i32) -> Self {
Self {
error: None,
result: ExecResult {
stdout: stdout.to_string(),
stderr: String::new(),
exit_code,
duration: Duration::from_millis(42),
},
}
}
fn failing(err: SandboxError) -> Self {
Self {
error: Some(err),
result: ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: 0,
duration: Duration::ZERO,
},
}
}
}
#[async_trait]
impl SandboxBackend for MockBackend {
fn name(&self) -> &str {
"mock"
}
fn capabilities(&self) -> BackendCapabilities {
BackendCapabilities {
supported_languages: vec![Language::Python],
isolation_class: "mock".to_string(),
enforced_limits: EnforcedLimits {
timeout: true,
memory: false,
network_isolation: false,
filesystem_isolation: false,
environment_isolation: false,
},
}
}
async fn execute(&self, _request: ExecRequest) -> Result<ExecResult, SandboxError> {
if let Some(ref err) = self.error { Err(err.clone()) } else { Ok(self.result.clone()) }
}
}
struct MockToolContext {
content: Content,
actions: Mutex<EventActions>,
}
impl MockToolContext {
fn new() -> Self {
Self { content: Content::new("user"), actions: Mutex::new(EventActions::default()) }
}
}
#[async_trait]
impl ReadonlyContext for MockToolContext {
fn invocation_id(&self) -> &str {
"inv-1"
}
fn agent_name(&self) -> &str {
"test-agent"
}
fn user_id(&self) -> &str {
"user"
}
fn app_name(&self) -> &str {
"app"
}
fn session_id(&self) -> &str {
"session"
}
fn branch(&self) -> &str {
""
}
fn user_content(&self) -> &Content {
&self.content
}
}
#[async_trait]
impl CallbackContext for MockToolContext {
fn artifacts(&self) -> Option<Arc<dyn adk_core::Artifacts>> {
None
}
}
#[async_trait]
impl ToolContext for MockToolContext {
fn function_call_id(&self) -> &str {
"call-1"
}
fn actions(&self) -> EventActions {
self.actions.lock().unwrap().clone()
}
fn set_actions(&self, actions: EventActions) {
*self.actions.lock().unwrap() = actions;
}
async fn search_memory(
&self,
_query: &str,
) -> adk_core::Result<Vec<adk_core::MemoryEntry>> {
Ok(vec![])
}
}
fn ctx() -> Arc<dyn ToolContext> {
Arc::new(MockToolContext::new())
}
#[test]
fn test_name() {
let tool = SandboxTool::new(Arc::new(MockBackend::success("", 0)));
assert_eq!(tool.name(), "sandbox_exec");
}
#[test]
fn test_required_scopes() {
let tool = SandboxTool::new(Arc::new(MockBackend::success("", 0)));
assert_eq!(tool.required_scopes(), &["code:execute"]);
}
#[test]
fn test_parameters_schema_is_valid() {
let tool = SandboxTool::new(Arc::new(MockBackend::success("", 0)));
let schema = tool.parameters_schema().expect("schema should be Some");
assert_eq!(schema["type"], "object");
assert!(schema["properties"]["language"].is_object());
assert!(schema["properties"]["code"].is_object());
assert!(schema["properties"]["stdin"].is_object());
assert!(schema["properties"]["timeout_secs"].is_object());
let required = schema["required"].as_array().unwrap();
let required_strs: Vec<&str> = required.iter().map(|v| v.as_str().unwrap()).collect();
assert!(required_strs.contains(&"language"));
assert!(required_strs.contains(&"code"));
assert!(!required_strs.contains(&"stdin"));
assert!(!required_strs.contains(&"timeout_secs"));
}
#[tokio::test]
async fn test_successful_execution() {
let backend = Arc::new(MockBackend::success("hello\n", 0));
let tool = SandboxTool::new(backend);
let args = json!({ "language": "python", "code": "print('hello')" });
let result = tool.execute(ctx(), args).await.unwrap();
assert_eq!(result["status"], "success");
assert_eq!(result["stdout"], "hello\n");
assert_eq!(result["exit_code"], 0);
assert!(result["duration_ms"].is_number());
}
#[tokio::test]
async fn test_timeout_error_as_information() {
let backend = Arc::new(MockBackend::failing(SandboxError::Timeout {
timeout: Duration::from_secs(5),
}));
let tool = SandboxTool::new(backend);
let args = json!({ "language": "python", "code": "import time; time.sleep(100)" });
let result = tool.execute(ctx(), args).await;
assert!(result.is_ok());
let val = result.unwrap();
assert_eq!(val["status"], "timeout");
assert!(val["stderr"].as_str().unwrap().contains("timed out"));
assert!(val["duration_ms"].is_number());
}
#[tokio::test]
async fn test_memory_exceeded_error_as_information() {
let backend = Arc::new(MockBackend::failing(SandboxError::MemoryExceeded { limit_mb: 64 }));
let tool = SandboxTool::new(backend);
let args = json!({ "language": "wasm", "code": "(module)" });
let result = tool.execute(ctx(), args).await.unwrap();
assert_eq!(result["status"], "memory_exceeded");
assert!(result["stderr"].as_str().unwrap().contains("64 MB"));
}
#[tokio::test]
async fn test_execution_failed_error_as_information() {
let backend = Arc::new(MockBackend::failing(SandboxError::ExecutionFailed(
"spawn failed".to_string(),
)));
let tool = SandboxTool::new(backend);
let args = json!({ "language": "python", "code": "x" });
let result = tool.execute(ctx(), args).await.unwrap();
assert_eq!(result["status"], "error");
assert_eq!(result["stderr"], "spawn failed");
}
#[tokio::test]
async fn test_missing_language_field() {
let backend = Arc::new(MockBackend::success("", 0));
let tool = SandboxTool::new(backend);
let args = json!({ "code": "print('hi')" });
let result = tool.execute(ctx(), args).await.unwrap();
assert_eq!(result["status"], "error");
assert!(result["stderr"].as_str().unwrap().contains("language"));
}
#[tokio::test]
async fn test_missing_code_field() {
let backend = Arc::new(MockBackend::success("", 0));
let tool = SandboxTool::new(backend);
let args = json!({ "language": "python" });
let result = tool.execute(ctx(), args).await.unwrap();
assert_eq!(result["status"], "error");
assert!(result["stderr"].as_str().unwrap().contains("code"));
}
#[tokio::test]
async fn test_unsupported_language() {
let backend = Arc::new(MockBackend::success("", 0));
let tool = SandboxTool::new(backend);
let args = json!({ "language": "cobol", "code": "DISPLAY 'HI'" });
let result = tool.execute(ctx(), args).await.unwrap();
assert_eq!(result["status"], "error");
assert!(result["stderr"].as_str().unwrap().contains("cobol"));
}
#[tokio::test]
async fn test_custom_timeout() {
let backend = Arc::new(MockBackend::success("ok", 0));
let tool = SandboxTool::new(backend);
let args = json!({ "language": "python", "code": "print('ok')", "timeout_secs": 60 });
let result = tool.execute(ctx(), args).await.unwrap();
assert_eq!(result["status"], "success");
}
#[tokio::test]
async fn test_stdin_passed_through() {
let backend = Arc::new(MockBackend::success("echo", 0));
let tool = SandboxTool::new(backend);
let args = json!({
"language": "python",
"code": "import sys; print(sys.stdin.read())",
"stdin": "hello"
});
let result = tool.execute(ctx(), args).await.unwrap();
assert_eq!(result["status"], "success");
}
#[test]
fn test_parse_language_all_variants() {
assert_eq!(parse_language(&json!("rust")).unwrap(), Language::Rust);
assert_eq!(parse_language(&json!("python")).unwrap(), Language::Python);
assert_eq!(parse_language(&json!("javascript")).unwrap(), Language::JavaScript);
assert_eq!(parse_language(&json!("typescript")).unwrap(), Language::TypeScript);
assert_eq!(parse_language(&json!("wasm")).unwrap(), Language::Wasm);
assert_eq!(parse_language(&json!("command")).unwrap(), Language::Command);
assert!(parse_language(&json!("ruby")).is_err());
assert!(parse_language(&json!(42)).is_err());
}
#[test]
fn test_sandbox_error_to_json_variants() {
let timeout_json =
sandbox_error_to_json(&SandboxError::Timeout { timeout: Duration::from_secs(10) });
assert_eq!(timeout_json["status"], "timeout");
let mem_json = sandbox_error_to_json(&SandboxError::MemoryExceeded { limit_mb: 128 });
assert_eq!(mem_json["status"], "memory_exceeded");
let exec_json = sandbox_error_to_json(&SandboxError::ExecutionFailed("boom".into()));
assert_eq!(exec_json["status"], "error");
let invalid_json = sandbox_error_to_json(&SandboxError::InvalidRequest("bad".into()));
assert_eq!(invalid_json["status"], "error");
let unavail_json = sandbox_error_to_json(&SandboxError::BackendUnavailable("gone".into()));
assert_eq!(unavail_json["status"], "error");
}
}