From 9a823447fa78d7d79bf7fa4c61e8d2efbcca654c Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Sat, 7 Mar 2026 23:13:51 -0500 Subject: [PATCH] Pass [sandbox.env] through API backend tool execution Previously sandbox env vars only reached CLI backend agents but not API backend tool calls. Thread tool_env through ToolContext so shell and web_fetch tools pass env vars to exec_command for all backends. Co-Authored-By: Claude Opus 4.6 --- crates/arc-agent/src/mcp_integration.rs | 1 + crates/arc-agent/src/session.rs | 7 ++ crates/arc-agent/src/skills.rs | 3 + crates/arc-agent/src/test_support.rs | 9 ++- crates/arc-agent/src/tool_execution.rs | 16 ++++ crates/arc-agent/src/tool_registry.rs | 2 + crates/arc-agent/src/tools.rs | 97 ++++++++++++++++++++++- crates/arc-workflows/src/cli/backend.rs | 25 +++++- crates/arc-workflows/src/cli/run.rs | 3 +- crates/arc-workflows/src/hook/executor.rs | 1 + 10 files changed, 157 insertions(+), 7 deletions(-) diff --git a/crates/arc-agent/src/mcp_integration.rs b/crates/arc-agent/src/mcp_integration.rs index 6740e7389..3df2c30e8 100644 --- a/crates/arc-agent/src/mcp_integration.rs +++ b/crates/arc-agent/src/mcp_integration.rs @@ -94,6 +94,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; diff --git a/crates/arc-agent/src/session.rs b/crates/arc-agent/src/session.rs index 51912d60d..4bb7c5649 100644 --- a/crates/arc-agent/src/session.rs +++ b/crates/arc-agent/src/session.rs @@ -40,6 +40,7 @@ pub struct Session { skills: Vec, system_prompt: String, file_tracker: FileTracker, + tool_env: Option>, } impl Session { @@ -67,9 +68,14 @@ impl Session { skills: Vec::new(), system_prompt: String::new(), file_tracker: FileTracker::default(), + tool_env: None, } } + pub fn set_tool_env(&mut self, env: std::collections::HashMap) { + self.tool_env = Some(env); + } + /// Initialize session by discovering project docs and capturing environment context. /// Call before `process_input`. pub async fn initialize(&mut self) { @@ -609,6 +615,7 @@ impl Session { &self.config, &self.event_emitter, &self.id, + self.tool_env.as_ref(), ) .await; diff --git a/crates/arc-agent/src/skills.rs b/crates/arc-agent/src/skills.rs index d8993a7ff..08a1195e2 100644 --- a/crates/arc-agent/src/skills.rs +++ b/crates/arc-agent/src/skills.rs @@ -548,6 +548,7 @@ name: trimmed let ctx = crate::tool_registry::ToolContext { env, cancel: tokio_util::sync::CancellationToken::new(), + tool_env: None, }; let result = (tool.executor)(args, ctx).await; assert_eq!( @@ -566,6 +567,7 @@ name: trimmed let ctx = crate::tool_registry::ToolContext { env, cancel: tokio_util::sync::CancellationToken::new(), + tool_env: None, }; let result = (tool.executor)(args, ctx).await; assert!(result.is_err()); @@ -582,6 +584,7 @@ name: trimmed let ctx = crate::tool_registry::ToolContext { env, cancel: tokio_util::sync::CancellationToken::new(), + tool_env: None, }; let result = (tool.executor)(args, ctx).await; assert!(result.is_err()); diff --git a/crates/arc-agent/src/test_support.rs b/crates/arc-agent/src/test_support.rs index 41d97097c..cd2d5ff3d 100644 --- a/crates/arc-agent/src/test_support.rs +++ b/crates/arc-agent/src/test_support.rs @@ -33,6 +33,8 @@ pub struct MockSandbox { pub captured_timeout: Mutex>, /// Captures the `command` argument from `exec_command` calls. pub captured_command: Mutex>, + /// Captures the `env_vars` argument from `exec_command` calls. + pub captured_env_vars: Mutex>>, pub event_callback: Option, } @@ -76,6 +78,7 @@ impl Default for MockSandbox { written_files: Mutex::new(Vec::new()), captured_timeout: Mutex::new(None), captured_command: Mutex::new(None), + captured_env_vars: Mutex::new(None), event_callback: None, } } @@ -135,7 +138,7 @@ impl Sandbox for MockSandbox { command: &str, timeout_ms: u64, _working_dir: Option<&str>, - _env_vars: Option<&std::collections::HashMap>, + env_vars: Option<&std::collections::HashMap>, _cancel_token: Option, ) -> Result { *self @@ -146,6 +149,10 @@ impl Sandbox for MockSandbox { .captured_command .lock() .expect("captured_command lock poisoned") = Some(command.to_string()); + *self + .captured_env_vars + .lock() + .expect("captured_env_vars lock poisoned") = env_vars.cloned(); Ok(self.exec_result.clone()) } diff --git a/crates/arc-agent/src/tool_execution.rs b/crates/arc-agent/src/tool_execution.rs index 23153a7e6..feb7bf7ab 100644 --- a/crates/arc-agent/src/tool_execution.rs +++ b/crates/arc-agent/src/tool_execution.rs @@ -5,6 +5,7 @@ use crate::tool_registry::ToolRegistry; use crate::truncation::truncate_tool_output; use crate::types::AgentEvent; use arc_llm::types::ToolResult; +use std::collections::HashMap; use std::sync::Arc; use tokio_util::sync::CancellationToken; @@ -20,6 +21,7 @@ pub async fn execute_tool_calls( config: &SessionConfig, emitter: &EventEmitter, session_id: &str, + tool_env: Option<&HashMap>, ) -> Vec { if parallel && tool_calls.len() > 1 { execute_tool_calls_parallel( @@ -31,6 +33,7 @@ pub async fn execute_tool_calls( config, emitter, session_id, + tool_env, ) .await } else { @@ -43,6 +46,7 @@ pub async fn execute_tool_calls( config, emitter, session_id, + tool_env, ) .await } @@ -58,6 +62,7 @@ async fn execute_tool_calls_sequential( config: &SessionConfig, emitter: &EventEmitter, session_id: &str, + tool_env: Option<&HashMap>, ) -> Vec { let mut results = Vec::new(); for tc in tool_calls { @@ -75,6 +80,7 @@ async fn execute_tool_calls_sequential( config, emitter, session_id, + tool_env, ) .await; results.push(result); @@ -92,7 +98,9 @@ async fn execute_tool_calls_parallel( config: &SessionConfig, emitter: &EventEmitter, session_id: &str, + tool_env: Option<&HashMap>, ) -> Vec { + let tool_env = tool_env.cloned(); let futures: Vec<_> = tool_calls .iter() .map(|tc| { @@ -103,6 +111,7 @@ async fn execute_tool_calls_parallel( let tc = tc.clone(); let session_id = session_id.to_owned(); let tool_approval = tool_approval.cloned(); + let tool_env = tool_env.clone(); // Look up the tool before spawning since ToolRegistry is not Send. let registered_tool = registry.get(&tc.name).cloned(); async move { @@ -115,6 +124,7 @@ async fn execute_tool_calls_parallel( &config, &emitter, &session_id, + tool_env.as_ref(), ) .await } @@ -135,6 +145,7 @@ pub async fn execute_and_emit_one_tool( config: &SessionConfig, emitter: &EventEmitter, session_id: &str, + tool_env: Option<&HashMap>, ) -> ToolResult { execute_and_emit_one_tool_with_lookup( tc, @@ -145,6 +156,7 @@ pub async fn execute_and_emit_one_tool( config, emitter, session_id, + tool_env, ) .await } @@ -160,6 +172,7 @@ async fn execute_and_emit_one_tool_with_lookup( config: &SessionConfig, emitter: &EventEmitter, session_id: &str, + tool_env: Option<&HashMap>, ) -> ToolResult { emitter.emit( session_id.to_owned(), @@ -178,6 +191,7 @@ async fn execute_and_emit_one_tool_with_lookup( env, tool_approval, cancel_token, + tool_env, ) .await; @@ -210,6 +224,7 @@ async fn execute_one_tool( env: Arc, tool_approval: Option<&ToolApprovalFn>, cancel_token: CancellationToken, + tool_env: Option<&HashMap>, ) -> ToolResult { if let Some(approval_fn) = tool_approval { if let Err(denial_message) = approval_fn(tool_name, arguments) { @@ -228,6 +243,7 @@ async fn execute_one_tool( let ctx = crate::tool_registry::ToolContext { env, cancel: cancel_token, + tool_env: tool_env.cloned(), }; match (tool.executor)(arguments.clone(), ctx).await { Ok(output) => ToolResult::success(tool_call_id, serde_json::json!(output)), diff --git a/crates/arc-agent/src/tool_registry.rs b/crates/arc-agent/src/tool_registry.rs index 6aedd6d56..cec5057a6 100644 --- a/crates/arc-agent/src/tool_registry.rs +++ b/crates/arc-agent/src/tool_registry.rs @@ -9,6 +9,7 @@ use tokio_util::sync::CancellationToken; pub struct ToolContext { pub env: Arc, pub cancel: CancellationToken, + pub tool_env: Option>, } pub type ToolExecutor = Arc< @@ -178,6 +179,7 @@ mod tests { let ctx = ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }; let result = (tool.executor)(serde_json::json!({}), ctx).await; assert_eq!(result.unwrap(), "ok"); diff --git a/crates/arc-agent/src/tools.rs b/crates/arc-agent/src/tools.rs index d46015504..6c5623637 100644 --- a/crates/arc-agent/src/tools.rs +++ b/crates/arc-agent/src/tools.rs @@ -219,9 +219,10 @@ pub fn make_shell_tool_with_config(config: &SessionConfig) -> RegisteredTool { .unwrap_or(default_timeout) .min(max_timeout); + tracing::debug!(env_var_count = ctx.tool_env.as_ref().map_or(0, |e| e.len()), "Injecting sandbox env vars into tool execution"); let result = ctx .env - .exec_command(command, timeout_ms, None, None, Some(ctx.cancel)) + .exec_command(command, timeout_ms, None, ctx.tool_env.as_ref(), Some(ctx.cancel)) .await?; let mut output = String::new(); @@ -534,7 +535,7 @@ pub(crate) fn make_web_fetch_tool(summarizer: Option) -> Reg ); let result = ctx.env - .exec_command(&command, timeout_ms, None, None, Some(ctx.cancel)) + .exec_command(&command, timeout_ms, None, ctx.tool_env.as_ref(), Some(ctx.cancel)) .await?; if result.exit_code != 0 { @@ -611,6 +612,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -635,6 +637,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -651,6 +654,7 @@ mod tests { ToolContext { env: env_clone, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -680,6 +684,7 @@ mod tests { ToolContext { env: env_clone, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -707,6 +712,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -731,6 +737,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -759,6 +766,7 @@ mod tests { ToolContext { env: env_clone, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -786,6 +794,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -804,6 +813,7 @@ mod tests { ToolContext { env: env_clone, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -828,6 +838,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -854,6 +865,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -861,6 +873,73 @@ mod tests { assert!(output.starts_with("Command timed out.\n")); } + #[tokio::test] + async fn shell_passes_tool_env_to_exec_command() { + let tool = make_shell_tool(); + let env = Arc::new(MockSandbox::default()); + let env_clone: Arc = env.clone(); + let mut tool_env = HashMap::new(); + tool_env.insert("MY_KEY".into(), "my_value".into()); + let _result = (tool.executor)( + serde_json::json!({"command": "echo $MY_KEY"}), + ToolContext { + env: env_clone, + cancel: CancellationToken::new(), + tool_env: Some(tool_env.clone()), + }, + ) + .await; + let captured = env.captured_env_vars.lock().unwrap().clone(); + assert_eq!(captured, Some(tool_env)); + } + + #[tokio::test] + async fn shell_passes_none_env_when_tool_env_is_none() { + let tool = make_shell_tool(); + let env = Arc::new(MockSandbox::default()); + let env_clone: Arc = env.clone(); + let _result = (tool.executor)( + serde_json::json!({"command": "echo hello"}), + ToolContext { + env: env_clone, + cancel: CancellationToken::new(), + tool_env: None, + }, + ) + .await; + let captured = env.captured_env_vars.lock().unwrap().clone(); + assert_eq!(captured, None); + } + + #[tokio::test] + async fn web_fetch_passes_tool_env_to_exec_command() { + let tool = make_web_fetch_tool(None); + let env = Arc::new(MockSandbox { + exec_result: ExecResult { + stdout: "fetched content".into(), + stderr: String::new(), + exit_code: 0, + timed_out: false, + duration_ms: 100, + }, + ..Default::default() + }); + let env_clone: Arc = env.clone(); + let mut tool_env = HashMap::new(); + tool_env.insert("API_KEY".into(), "secret".into()); + let _result = (tool.executor)( + serde_json::json!({"url": "https://example.com"}), + ToolContext { + env: env_clone, + cancel: CancellationToken::new(), + tool_env: Some(tool_env.clone()), + }, + ) + .await; + let captured = env.captured_env_vars.lock().unwrap().clone(); + assert_eq!(captured, Some(tool_env)); + } + #[tokio::test] async fn grep_basic() { let tool = make_grep_tool(); @@ -876,6 +955,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -896,6 +976,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -913,6 +994,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -932,6 +1014,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -984,6 +1067,7 @@ mod tests { ToolContext { env: env_clone, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -1020,6 +1104,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -1040,6 +1125,7 @@ mod tests { ToolContext { env: env_clone, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -1061,6 +1147,7 @@ mod tests { ToolContext { env: env_clone, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -1091,6 +1178,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -1117,6 +1205,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -1160,6 +1249,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -1189,6 +1279,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -1251,6 +1342,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; @@ -1301,6 +1393,7 @@ mod tests { ToolContext { env, cancel: CancellationToken::new(), + tool_env: None, }, ) .await; diff --git a/crates/arc-workflows/src/cli/backend.rs b/crates/arc-workflows/src/cli/backend.rs index 6a27d5c46..e9343c055 100644 --- a/crates/arc-workflows/src/cli/backend.rs +++ b/crates/arc-workflows/src/cli/backend.rs @@ -107,6 +107,7 @@ pub struct AgentApiBackend { provider: Provider, fallback_chain: Vec, sessions: Mutex>, + env: HashMap, } impl AgentApiBackend { @@ -117,9 +118,16 @@ impl AgentApiBackend { provider, fallback_chain, sessions: Mutex::new(HashMap::new()), + env: HashMap::new(), } } + #[must_use] + pub fn with_env(mut self, env: HashMap) -> Self { + self.env = env; + self + } + async fn create_session( &self, node: &Node, @@ -130,6 +138,7 @@ impl AgentApiBackend { self.provider, node, sandbox, + &self.env, ) .await } @@ -139,6 +148,7 @@ impl AgentApiBackend { provider: Provider, node: &Node, sandbox: &Arc, + env: &HashMap, ) -> Result { let client = Client::from_env() .await @@ -161,6 +171,7 @@ impl AgentApiBackend { let factory_client = client.clone(); let factory_model = model.to_string(); let factory_env = Arc::clone(sandbox); + let factory_tool_env = env.clone(); let factory: SessionFactory = Arc::new(move || { let child_profile: Arc = match provider { Provider::OpenAi => Arc::new(OpenAiProfile::new(&factory_model)), @@ -170,18 +181,25 @@ impl AgentApiBackend { Provider::Gemini => Arc::new(GeminiProfile::new(&factory_model)), Provider::Anthropic => Arc::new(AnthropicProfile::new(&factory_model)), }; - Session::new( + let mut session = Session::new( factory_client.clone(), child_profile, Arc::clone(&factory_env), SessionConfig::default(), - ) + ); + if !factory_tool_env.is_empty() { + session.set_tool_env(factory_tool_env.clone()); + } + session }); profile.register_subagent_tools(manager, factory, 0); let profile: Arc = Arc::from(profile); - let session = Session::new(client, profile, Arc::clone(sandbox), config); + let mut session = Session::new(client, profile, Arc::clone(sandbox), config); + if !env.is_empty() { + session.set_tool_env(env.clone()); + } // Wire subagent event callback to parent session's emitter manager_for_callback @@ -441,6 +459,7 @@ impl CodergenBackend for AgentApiBackend { target_provider, node, sandbox, + &self.env, ) .await { diff --git a/crates/arc-workflows/src/cli/run.rs b/crates/arc-workflows/src/cli/run.rs index 4c2cf6a16..11b7ca04f 100644 --- a/crates/arc-workflows/src/cli/run.rs +++ b/crates/arc-workflows/src/cli/run.rs @@ -671,7 +671,8 @@ pub async fn run_command( if dry_run_mode { None } else { - let api = AgentApiBackend::new(model.clone(), provider_enum, fallback_chain.clone()); + let api = AgentApiBackend::new(model.clone(), provider_enum, fallback_chain.clone()) + .with_env(sandbox_env.clone()); let cli = AgentCliBackend::new(model.clone(), provider_enum) .with_env(sandbox_env.clone()); Some(Box::new(BackendRouter::new(Box::new(api), cli))) diff --git a/crates/arc-workflows/src/hook/executor.rs b/crates/arc-workflows/src/hook/executor.rs index 45df35cb4..d9ed83ab2 100644 --- a/crates/arc-workflows/src/hook/executor.rs +++ b/crates/arc-workflows/src/hook/executor.rs @@ -365,6 +365,7 @@ impl HookExecutorImpl { let ctx = arc_agent::tool_registry::ToolContext { env: sandbox.clone(), cancel: cancel.child_token(), + tool_env: None, }; let result = match tool { Some(t) => match (t.executor)(tc.arguments.clone(), ctx).await {