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 <noreply@anthropic.com>
This commit is contained in:
Bryan Helmkamp 2026-03-07 23:13:51 -05:00
parent 13ee1735c0
commit 9a823447fa
10 changed files with 157 additions and 7 deletions

View file

@ -94,6 +94,7 @@ mod tests {
ToolContext {
env,
cancel: CancellationToken::new(),
tool_env: None,
},
)
.await;

View file

@ -40,6 +40,7 @@ pub struct Session {
skills: Vec<Skill>,
system_prompt: String,
file_tracker: FileTracker,
tool_env: Option<std::collections::HashMap<String, String>>,
}
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<String, String>) {
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;

View file

@ -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());

View file

@ -33,6 +33,8 @@ pub struct MockSandbox {
pub captured_timeout: Mutex<Option<u64>>,
/// Captures the `command` argument from `exec_command` calls.
pub captured_command: Mutex<Option<String>>,
/// Captures the `env_vars` argument from `exec_command` calls.
pub captured_env_vars: Mutex<Option<HashMap<String, String>>>,
pub event_callback: Option<crate::sandbox::SandboxEventCallback>,
}
@ -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<String, String>>,
env_vars: Option<&std::collections::HashMap<String, String>>,
_cancel_token: Option<CancellationToken>,
) -> Result<ExecResult, String> {
*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())
}

View file

@ -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<String, String>>,
) -> Vec<ToolResult> {
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<String, String>>,
) -> Vec<ToolResult> {
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<String, String>>,
) -> Vec<ToolResult> {
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<String, String>>,
) -> 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<String, String>>,
) -> 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<dyn Sandbox>,
tool_approval: Option<&ToolApprovalFn>,
cancel_token: CancellationToken,
tool_env: Option<&HashMap<String, String>>,
) -> 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)),

View file

@ -9,6 +9,7 @@ use tokio_util::sync::CancellationToken;
pub struct ToolContext {
pub env: Arc<dyn Sandbox>,
pub cancel: CancellationToken,
pub tool_env: Option<HashMap<String, String>>,
}
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");

View file

@ -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<WebFetchSummarizer>) -> 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<dyn Sandbox> = 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<dyn Sandbox> = 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<dyn Sandbox> = 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;

View file

@ -107,6 +107,7 @@ pub struct AgentApiBackend {
provider: Provider,
fallback_chain: Vec<FallbackTarget>,
sessions: Mutex<HashMap<String, Session>>,
env: HashMap<String, String>,
}
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<String, String>) -> 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<dyn Sandbox>,
env: &HashMap<String, String>,
) -> Result<Session, ArcError> {
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<dyn ProviderProfile> = 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<dyn ProviderProfile> = 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
{

View file

@ -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)))

View file

@ -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 {