use std::collections::BTreeSet;
use std::sync::Arc;
use async_trait::async_trait;
use serde_json::Value;
use adk_agent::LlmAgentBuilder;
use adk_core::{Agent, Llm, Tool, ToolConfirmationPolicy, ToolContext};
#[cfg(feature = "sandbox")]
use adk_sandbox::{ExecRequest, Language, SandboxBackend};
use crate::types::{ManagedAgentDef, PermissionMode, PermissionPolicy, ToolConfig};
#[derive(Debug, thiserror::Error)]
pub enum BuildError {
#[error("invalid agent definition: {0}")]
InvalidDef(String),
#[error("agent build failed: {0}")]
BuildFailed(String),
}
#[cfg(feature = "sandbox")]
pub fn build_agent(
def: &ManagedAgentDef,
model: Arc<dyn Llm>,
sandbox: Option<Arc<dyn SandboxBackend>>,
) -> Result<Arc<dyn Agent>, BuildError> {
let mut builder = LlmAgentBuilder::new(&def.name).model(model);
if let Some(ref system) = def.system {
builder = builder.instruction(system.clone());
}
if let Some(ref description) = def.description {
builder = builder.description(description.clone());
}
for tool_config in &def.tools {
let tool: Arc<dyn Tool> = match tool_config {
ToolConfig::Bash {} => Arc::new(ManagedBuiltinTool::new(
"bash",
"Execute bash shell commands in the agent's workspace.",
sandbox.clone(),
)),
ToolConfig::Filesystem {} => Arc::new(ManagedBuiltinTool::new(
"filesystem",
"Read, write, and manage files in the agent's workspace.",
sandbox.clone(),
)),
ToolConfig::WebSearch {} => Arc::new(ManagedBuiltinTool::new(
"web_search",
"Search the web for information.",
sandbox.clone(),
)),
ToolConfig::WebFetch {} => Arc::new(ManagedBuiltinTool::new(
"web_fetch",
"Fetch and extract content from a URL.",
sandbox.clone(),
)),
ToolConfig::CodeExecution {} => Arc::new(ManagedBuiltinTool::new(
"code_execution",
"Execute code in a sandboxed environment.",
sandbox.clone(),
)),
ToolConfig::Custom { name, description, input_schema } => {
Arc::new(ManagedCustomTool::new(
name.clone(),
description.clone().unwrap_or_default(),
input_schema.clone(),
))
}
};
builder = builder.tool(tool);
}
if let Some(ref policy) = def.permission_policy {
let confirmation_policy = map_permission_policy(policy);
builder = builder.tool_confirmation_policy(confirmation_policy);
}
if !def.mcp_servers.is_empty() {
tracing::debug!(
mcp_count = def.mcp_servers.len(),
"MCP server configs noted (wiring deferred to session loop)"
);
}
if !def.skills.is_empty() {
tracing::debug!(
skill_count = def.skills.len(),
"skill refs noted (wiring deferred to session loop)"
);
}
let agent = builder.build().map_err(|e| BuildError::BuildFailed(e.to_string()))?;
Ok(Arc::new(agent))
}
#[cfg(not(feature = "sandbox"))]
pub fn build_agent(
def: &ManagedAgentDef,
model: Arc<dyn Llm>,
) -> Result<Arc<dyn Agent>, BuildError> {
let mut builder = LlmAgentBuilder::new(&def.name).model(model);
if let Some(ref system) = def.system {
builder = builder.instruction(system.clone());
}
if let Some(ref description) = def.description {
builder = builder.description(description.clone());
}
for tool_config in &def.tools {
let tool: Arc<dyn Tool> = match tool_config {
ToolConfig::Bash {} => Arc::new(ManagedBuiltinTool::new(
"bash",
"Execute bash shell commands in the agent's workspace.",
)),
ToolConfig::Filesystem {} => Arc::new(ManagedBuiltinTool::new(
"filesystem",
"Read, write, and manage files in the agent's workspace.",
)),
ToolConfig::WebSearch {} => {
Arc::new(ManagedBuiltinTool::new("web_search", "Search the web for information."))
}
ToolConfig::WebFetch {} => Arc::new(ManagedBuiltinTool::new(
"web_fetch",
"Fetch and extract content from a URL.",
)),
ToolConfig::CodeExecution {} => Arc::new(ManagedBuiltinTool::new(
"code_execution",
"Execute code in a sandboxed environment.",
)),
ToolConfig::Custom { name, description, input_schema } => {
Arc::new(ManagedCustomTool::new(
name.clone(),
description.clone().unwrap_or_default(),
input_schema.clone(),
))
}
};
builder = builder.tool(tool);
}
if let Some(ref policy) = def.permission_policy {
let confirmation_policy = map_permission_policy(policy);
builder = builder.tool_confirmation_policy(confirmation_policy);
}
if !def.mcp_servers.is_empty() {
tracing::debug!(
mcp_count = def.mcp_servers.len(),
"MCP server configs noted (wiring deferred to session loop)"
);
}
if !def.skills.is_empty() {
tracing::debug!(
skill_count = def.skills.len(),
"skill refs noted (wiring deferred to session loop)"
);
}
let agent = builder.build().map_err(|e| BuildError::BuildFailed(e.to_string()))?;
Ok(Arc::new(agent))
}
fn map_permission_policy(policy: &PermissionPolicy) -> ToolConfirmationPolicy {
let tools_requiring_confirmation: BTreeSet<String> = policy
.tools
.iter()
.filter(|(_, mode)| matches!(mode, PermissionMode::Prompt | PermissionMode::Deny))
.map(|(name, _)| name.clone())
.collect();
match policy.default {
PermissionMode::AutoApprove => {
if tools_requiring_confirmation.is_empty() {
ToolConfirmationPolicy::Never
} else {
ToolConfirmationPolicy::PerTool(tools_requiring_confirmation)
}
}
PermissionMode::Prompt | PermissionMode::Deny => {
ToolConfirmationPolicy::Always
}
}
}
#[derive(Clone)]
pub struct ManagedBuiltinTool {
name: String,
description: String,
#[cfg(feature = "sandbox")]
sandbox: Option<Arc<dyn SandboxBackend>>,
}
impl std::fmt::Debug for ManagedBuiltinTool {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut d = f.debug_struct("ManagedBuiltinTool");
d.field("name", &self.name).field("description", &self.description);
#[cfg(feature = "sandbox")]
d.field("sandbox", &self.sandbox.as_ref().map(|s| s.name()));
d.finish()
}
}
impl ManagedBuiltinTool {
#[cfg(feature = "sandbox")]
pub fn new(
name: impl Into<String>,
description: impl Into<String>,
sandbox: Option<Arc<dyn SandboxBackend>>,
) -> Self {
Self { name: name.into(), description: description.into(), sandbox }
}
#[cfg(not(feature = "sandbox"))]
pub fn new(name: impl Into<String>, description: impl Into<String>) -> Self {
Self { name: name.into(), description: description.into() }
}
async fn execute_bash(&self, args: &Value) -> adk_core::Result<Value> {
let command = args.get("command").and_then(|v| v.as_str()).unwrap_or_default();
if command.is_empty() {
return Ok(serde_json::json!({
"error": "no command provided",
"exit_code": 1
}));
}
let output = tokio::process::Command::new("sh")
.arg("-c")
.arg(command)
.output()
.await
.map_err(|e| adk_core::AdkError::tool(format!("failed to spawn bash: {e}")))?;
let stdout = String::from_utf8_lossy(&output.stdout).to_string();
let stderr = String::from_utf8_lossy(&output.stderr).to_string();
let exit_code = output.status.code().unwrap_or(-1);
Ok(serde_json::json!({
"stdout": stdout,
"stderr": stderr,
"exit_code": exit_code
}))
}
async fn execute_filesystem(&self, args: &Value) -> adk_core::Result<Value> {
let operation = args.get("operation").and_then(|v| v.as_str()).unwrap_or("read");
let path = args.get("path").and_then(|v| v.as_str()).unwrap_or("");
match operation {
"read" => {
if path.is_empty() {
return Ok(serde_json::json!({"error": "path is required"}));
}
match tokio::fs::read_to_string(path).await {
Ok(content) => Ok(serde_json::json!({"content": content})),
Err(e) => Ok(serde_json::json!({"error": format!("read failed: {e}")})),
}
}
"write" => {
let content = args.get("content").and_then(|v| v.as_str()).unwrap_or("");
if path.is_empty() {
return Ok(serde_json::json!({"error": "path is required"}));
}
match tokio::fs::write(path, content).await {
Ok(()) => Ok(serde_json::json!({"status": "written", "path": path})),
Err(e) => Ok(serde_json::json!({"error": format!("write failed: {e}")})),
}
}
"list" => {
let target = if path.is_empty() { "." } else { path };
match tokio::fs::read_dir(target).await {
Ok(mut entries) => {
let mut files = Vec::new();
while let Ok(Some(entry)) = entries.next_entry().await {
files.push(entry.file_name().to_string_lossy().to_string());
}
Ok(serde_json::json!({"files": files}))
}
Err(e) => Ok(serde_json::json!({"error": format!("list failed: {e}")})),
}
}
other => Ok(serde_json::json!({
"error": format!("unsupported filesystem operation: {other}")
})),
}
}
#[cfg(feature = "sandbox")]
async fn execute_via_sandbox(
&self,
sandbox: &Arc<dyn SandboxBackend>,
language: Language,
args: &Value,
) -> adk_core::Result<Value> {
use std::collections::HashMap;
use std::time::Duration;
let code = match language {
Language::Command => {
args.get("command").and_then(|v| v.as_str()).unwrap_or_default().to_string()
}
_ => {
args.get("code").and_then(|v| v.as_str()).unwrap_or_default().to_string()
}
};
if code.is_empty() {
return Ok(serde_json::json!({"error": "no code/command provided"}));
}
let timeout_secs = args.get("timeout").and_then(|v| v.as_u64()).unwrap_or(30);
let request = ExecRequest {
language,
code,
stdin: args.get("stdin").and_then(|v| v.as_str()).map(String::from),
timeout: Duration::from_secs(timeout_secs),
memory_limit_mb: args.get("memory_limit_mb").and_then(|v| v.as_u64()).map(|v| v as u32),
env: HashMap::new(),
};
match sandbox.execute(request).await {
Ok(result) => Ok(serde_json::json!({
"stdout": result.stdout,
"stderr": result.stderr,
"exit_code": result.exit_code,
"duration_ms": result.duration.as_millis() as u64
})),
Err(e) => Ok(serde_json::json!({
"error": format!("sandbox execution failed: {e}"),
"exit_code": -1
})),
}
}
}
#[async_trait]
impl Tool for ManagedBuiltinTool {
fn name(&self) -> &str {
&self.name
}
fn description(&self) -> &str {
&self.description
}
async fn execute(&self, _ctx: Arc<dyn ToolContext>, args: Value) -> adk_core::Result<Value> {
match self.name.as_str() {
"bash" => {
#[cfg(feature = "sandbox")]
if let Some(ref sandbox) = self.sandbox {
return self.execute_via_sandbox(sandbox, Language::Command, &args).await;
}
self.execute_bash(&args).await
}
"filesystem" => self.execute_filesystem(&args).await,
"code_execution" => {
let language = args.get("language").and_then(|v| v.as_str()).unwrap_or("python");
let code = args.get("code").and_then(|v| v.as_str()).unwrap_or_default();
if code.is_empty() {
return Ok(serde_json::json!({"error": "no code provided"}));
}
#[cfg(feature = "sandbox")]
if let Some(ref sandbox) = self.sandbox {
let lang = match language {
"python" | "python3" => Language::Python,
"javascript" | "js" | "node" => Language::JavaScript,
"bash" | "sh" => Language::Command,
"rust" => Language::Rust,
"typescript" | "ts" => Language::TypeScript,
other => {
return Ok(serde_json::json!({
"error": format!("unsupported language for sandbox: {other}")
}));
}
};
return self.execute_via_sandbox(sandbox, lang, &args).await;
}
let interpreter = match language {
"python" | "python3" => "python3",
"javascript" | "js" | "node" => "node",
"bash" | "sh" => "sh",
other => {
return Ok(serde_json::json!({
"error": format!("unsupported language: {other}. Configure a sandbox backend for full language support.")
}));
}
};
let output = tokio::process::Command::new(interpreter)
.arg("-c")
.arg(code)
.output()
.await
.map_err(|e| {
adk_core::AdkError::tool(format!("failed to spawn {interpreter}: {e}"))
})?;
let stdout = String::from_utf8_lossy(&output.stdout).to_string();
let stderr = String::from_utf8_lossy(&output.stderr).to_string();
let exit_code = output.status.code().unwrap_or(-1);
Ok(serde_json::json!({
"stdout": stdout,
"stderr": stderr,
"exit_code": exit_code
}))
}
"web_search" => {
let query = args.get("query").and_then(|v| v.as_str()).unwrap_or_default();
Ok(serde_json::json!({
"error": "web_search is not configured for in-process execution. Configure an API key or use a provider with built-in search grounding.",
"query": query
}))
}
"web_fetch" => {
let url = args.get("url").and_then(|v| v.as_str()).unwrap_or_default();
Ok(serde_json::json!({
"error": "web_fetch is not configured for in-process execution. Configure an HTTP client or sandbox with network access.",
"url": url
}))
}
other => Err(adk_core::AdkError::tool(format!("unknown built-in tool: {other}"))),
}
}
}
#[derive(Debug, Clone)]
pub struct ManagedCustomTool {
name: String,
description: String,
input_schema: Value,
}
impl ManagedCustomTool {
pub fn new(name: String, description: String, input_schema: Value) -> Self {
Self { name, description, input_schema }
}
}
#[async_trait]
impl Tool for ManagedCustomTool {
fn name(&self) -> &str {
&self.name
}
fn description(&self) -> &str {
&self.description
}
fn parameters_schema(&self) -> Option<Value> {
Some(self.input_schema.clone())
}
fn is_long_running(&self) -> bool {
true
}
async fn execute(&self, _ctx: Arc<dyn ToolContext>, args: Value) -> adk_core::Result<Value> {
Ok(serde_json::json!({
"status": "pending_client_execution",
"tool": self.name,
"message": "This tool requires client-side execution. The result will be provided by the client.",
"args_received": args
}))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::{ManagedAgentDef, ModelRef, PermissionMode, PermissionPolicy, ToolConfig};
use adk_core::{Content, FinishReason, Llm, LlmRequest, LlmResponse, LlmResponseStream};
use async_stream::stream;
use std::collections::HashMap;
struct MockLlm {
name: String,
}
impl MockLlm {
fn new(name: &str) -> Self {
Self { name: name.to_string() }
}
}
#[async_trait]
impl Llm for MockLlm {
fn name(&self) -> &str {
&self.name
}
async fn generate_content(
&self,
_request: LlmRequest,
_stream: bool,
) -> adk_core::Result<LlmResponseStream> {
let s = stream! {
yield Ok(LlmResponse {
content: Some(Content::new("model").with_text("Hello")),
partial: false,
turn_complete: true,
finish_reason: Some(FinishReason::Stop),
..Default::default()
});
};
Ok(Box::pin(s))
}
}
#[cfg(feature = "sandbox")]
fn test_build_agent(def: &ManagedAgentDef, model: Arc<dyn Llm>) -> Arc<dyn Agent> {
build_agent(def, model, None).unwrap()
}
#[cfg(not(feature = "sandbox"))]
fn test_build_agent(def: &ManagedAgentDef, model: Arc<dyn Llm>) -> Arc<dyn Agent> {
build_agent(def, model).unwrap()
}
#[test]
fn test_build_agent_minimal_def() {
let def = ManagedAgentDef {
name: "test-agent".to_string(),
model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
system: None,
description: None,
tools: vec![],
mcp_servers: vec![],
skills: vec![],
permission_policy: None,
metadata: None,
};
let model: Arc<dyn Llm> = Arc::new(MockLlm::new("mock-model"));
let agent = test_build_agent(&def, model);
assert_eq!(agent.name(), "test-agent");
}
#[test]
fn test_build_agent_with_system_prompt() {
let def = ManagedAgentDef {
name: "prompted-agent".to_string(),
model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
system: Some("You are a helpful assistant.".to_string()),
description: Some("A helpful agent".to_string()),
tools: vec![],
mcp_servers: vec![],
skills: vec![],
permission_policy: None,
metadata: None,
};
let model: Arc<dyn Llm> = Arc::new(MockLlm::new("mock-model"));
let agent = test_build_agent(&def, model);
assert_eq!(agent.name(), "prompted-agent");
assert_eq!(agent.description(), "A helpful agent");
}
#[test]
fn test_build_agent_with_builtin_tools() {
let def = ManagedAgentDef {
name: "tool-agent".to_string(),
model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
system: None,
description: None,
tools: vec![
ToolConfig::Bash {},
ToolConfig::Filesystem {},
ToolConfig::WebSearch {},
ToolConfig::WebFetch {},
ToolConfig::CodeExecution {},
],
mcp_servers: vec![],
skills: vec![],
permission_policy: None,
metadata: None,
};
let model: Arc<dyn Llm> = Arc::new(MockLlm::new("mock-model"));
let agent = test_build_agent(&def, model);
assert_eq!(agent.name(), "tool-agent");
}
#[test]
fn test_build_agent_with_custom_tool() {
let def = ManagedAgentDef {
name: "custom-tool-agent".to_string(),
model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
system: None,
description: None,
tools: vec![ToolConfig::Custom {
name: "get_weather".to_string(),
description: Some("Get current weather".to_string()),
input_schema: serde_json::json!({
"type": "object",
"properties": {
"city": {"type": "string"}
},
"required": ["city"]
}),
}],
mcp_servers: vec![],
skills: vec![],
permission_policy: None,
metadata: None,
};
let model: Arc<dyn Llm> = Arc::new(MockLlm::new("mock-model"));
let agent = test_build_agent(&def, model);
assert_eq!(agent.name(), "custom-tool-agent");
}
#[test]
fn test_build_agent_with_permission_policy_auto_approve() {
let def = ManagedAgentDef {
name: "auto-agent".to_string(),
model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
system: None,
description: None,
tools: vec![ToolConfig::Bash {}],
mcp_servers: vec![],
skills: vec![],
permission_policy: Some(PermissionPolicy {
default: PermissionMode::AutoApprove,
tools: HashMap::new(),
}),
metadata: None,
};
let model: Arc<dyn Llm> = Arc::new(MockLlm::new("mock-model"));
let agent = test_build_agent(&def, model);
assert_eq!(agent.name(), "auto-agent");
}
#[test]
fn test_build_agent_with_permission_policy_prompt_default() {
let def = ManagedAgentDef {
name: "prompt-agent".to_string(),
model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
system: None,
description: None,
tools: vec![ToolConfig::Bash {}],
mcp_servers: vec![],
skills: vec![],
permission_policy: Some(PermissionPolicy {
default: PermissionMode::Prompt,
tools: HashMap::new(),
}),
metadata: None,
};
let model: Arc<dyn Llm> = Arc::new(MockLlm::new("mock-model"));
let agent = test_build_agent(&def, model);
assert_eq!(agent.name(), "prompt-agent");
}
#[test]
fn test_build_agent_with_per_tool_permission() {
let def = ManagedAgentDef {
name: "mixed-agent".to_string(),
model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
system: None,
description: None,
tools: vec![ToolConfig::Bash {}, ToolConfig::Filesystem {}],
mcp_servers: vec![],
skills: vec![],
permission_policy: Some(PermissionPolicy {
default: PermissionMode::AutoApprove,
tools: HashMap::from([
("bash".to_string(), PermissionMode::Prompt),
("delete_file".to_string(), PermissionMode::Deny),
]),
}),
metadata: None,
};
let model: Arc<dyn Llm> = Arc::new(MockLlm::new("mock-model"));
let agent = test_build_agent(&def, model);
assert_eq!(agent.name(), "mixed-agent");
}
#[test]
fn test_map_auto_approve_no_overrides() {
let policy =
PermissionPolicy { default: PermissionMode::AutoApprove, tools: HashMap::new() };
assert_eq!(map_permission_policy(&policy), ToolConfirmationPolicy::Never);
}
#[test]
fn test_map_prompt_default() {
let policy = PermissionPolicy { default: PermissionMode::Prompt, tools: HashMap::new() };
assert_eq!(map_permission_policy(&policy), ToolConfirmationPolicy::Always);
}
#[test]
fn test_map_deny_default() {
let policy = PermissionPolicy { default: PermissionMode::Deny, tools: HashMap::new() };
assert_eq!(map_permission_policy(&policy), ToolConfirmationPolicy::Always);
}
#[test]
fn test_map_auto_approve_with_per_tool_prompt() {
let policy = PermissionPolicy {
default: PermissionMode::AutoApprove,
tools: HashMap::from([
("bash".to_string(), PermissionMode::Prompt),
("delete_file".to_string(), PermissionMode::Deny),
]),
};
let result = map_permission_policy(&policy);
match result {
ToolConfirmationPolicy::PerTool(tools) => {
assert!(tools.contains("bash"));
assert!(tools.contains("delete_file"));
assert_eq!(tools.len(), 2);
}
other => panic!("expected PerTool, got: {other:?}"),
}
}
#[test]
fn test_map_auto_approve_with_auto_approve_overrides_only() {
let policy = PermissionPolicy {
default: PermissionMode::AutoApprove,
tools: HashMap::from([("read_file".to_string(), PermissionMode::AutoApprove)]),
};
assert_eq!(map_permission_policy(&policy), ToolConfirmationPolicy::Never);
}
#[cfg(feature = "sandbox")]
fn make_builtin_tool(name: &str, desc: &str) -> ManagedBuiltinTool {
ManagedBuiltinTool::new(name, desc, None)
}
#[cfg(not(feature = "sandbox"))]
fn make_builtin_tool(name: &str, desc: &str) -> ManagedBuiltinTool {
ManagedBuiltinTool::new(name, desc)
}
#[test]
fn test_builtin_tool_metadata() {
let tool = make_builtin_tool("bash", "Execute bash commands.");
assert_eq!(tool.name(), "bash");
assert_eq!(tool.description(), "Execute bash commands.");
}
#[tokio::test]
async fn test_builtin_tool_bash_executes() {
let tool = make_builtin_tool("bash", "Execute bash commands.");
let ctx = Arc::new(adk_tool::SimpleToolContext::new("test-caller"));
let result = tool.execute(ctx, serde_json::json!({"command": "echo hello"})).await.unwrap();
assert_eq!(result["exit_code"], 0);
assert!(result["stdout"].as_str().unwrap().contains("hello"));
}
#[tokio::test]
async fn test_builtin_tool_web_search_returns_error() {
let tool = make_builtin_tool("web_search", "Search the web.");
let ctx = Arc::new(adk_tool::SimpleToolContext::new("test-caller"));
let result = tool.execute(ctx, serde_json::json!({"query": "rust lang"})).await.unwrap();
assert!(result["error"].as_str().unwrap().contains("not configured"));
}
#[test]
fn test_custom_tool_metadata() {
let schema = serde_json::json!({
"type": "object",
"properties": {"city": {"type": "string"}}
});
let tool = ManagedCustomTool::new(
"get_weather".to_string(),
"Get current weather".to_string(),
schema.clone(),
);
assert_eq!(tool.name(), "get_weather");
assert_eq!(tool.description(), "Get current weather");
assert_eq!(tool.parameters_schema(), Some(schema));
assert!(tool.is_long_running());
}
#[tokio::test]
async fn test_custom_tool_execute_returns_pending_status() {
let tool = ManagedCustomTool::new(
"my_tool".to_string(),
"A custom tool".to_string(),
serde_json::json!({"type": "object"}),
);
let ctx = Arc::new(adk_tool::SimpleToolContext::new("test-caller"));
let result = tool.execute(ctx, serde_json::json!({"city": "Seattle"})).await.unwrap();
assert_eq!(result["status"], "pending_client_execution");
assert_eq!(result["tool"], "my_tool");
assert_eq!(result["args_received"]["city"], "Seattle");
}
#[test]
fn test_custom_tool_is_long_running() {
let tool = ManagedCustomTool::new(
"deploy".to_string(),
"Deploy to production".to_string(),
serde_json::json!({"type": "object"}),
);
assert!(tool.is_long_running());
}
}