Skip to content

Commit dfaa580

Browse files
dingfelidingfeli
authored
fix(agent): tool permission override (#2606)
1 parent ccc6b33 commit dfaa580

File tree

7 files changed

+198
-68
lines changed

7 files changed

+198
-68
lines changed

‎crates/chat-cli/src/cli/agent/mod.rs‎

Lines changed: 79 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,10 @@ pub use wrapper_types::{
5454
tool_settings_schema,
5555
};
5656

57+
use super::chat::tools::execute::ExecuteCommand;
58+
use super::chat::tools::fs_read::FsRead;
59+
use super::chat::tools::fs_write::FsWrite;
60+
use super::chat::tools::use_aws::UseAws;
5761
use super::chat::tools::{
5862
DEFAULT_APPROVE,
5963
NATIVE_TOOLS,
@@ -213,8 +217,8 @@ impl Agent {
213217

214218
self.path = Some(path.to_path_buf());
215219

220+
let mut stderr = std::io::stderr();
216221
if let (true, Some(legacy_mcp_config)) = (self.use_legacy_mcp_json, legacy_mcp_config) {
217-
let mut stderr = std::io::stderr();
218222
for (name, legacy_server) in &legacy_mcp_config.mcp_servers {
219223
if mcp_servers.mcp_servers.contains_key(name) {
220224
let _ = queue!(
@@ -238,6 +242,58 @@ impl Agent {
238242
}
239243
}
240244

245+
stderr.flush()?;
246+
247+
Ok(())
248+
}
249+
250+
pub fn validate_tool_settings(&self, output: &mut impl Write) -> Result<(), AgentConfigError> {
251+
let execute_name = if cfg!(windows) { "execute_cmd" } else { "execute_bash" };
252+
for allowed_tool in &self.allowed_tools {
253+
if let Some(settings) = self.tools_settings.get(allowed_tool.as_str()) {
254+
// currently we only have four native tools that offers tool settings
255+
match allowed_tool.as_str() {
256+
"fs_read" => {
257+
if let Some(overridden_settings) = FsRead::allowable_field_to_be_overridden(settings) {
258+
queue_permission_override_warning(
259+
allowed_tool.as_str(),
260+
overridden_settings.as_str(),
261+
output,
262+
)?;
263+
}
264+
},
265+
"fs_write" => {
266+
if let Some(overridden_settings) = FsWrite::allowable_field_to_be_overridden(settings) {
267+
queue_permission_override_warning(
268+
allowed_tool.as_str(),
269+
overridden_settings.as_str(),
270+
output,
271+
)?;
272+
}
273+
},
274+
"use_aws" => {
275+
if let Some(overridden_settings) = UseAws::allowable_field_to_be_overridden(settings) {
276+
queue_permission_override_warning(
277+
allowed_tool.as_str(),
278+
overridden_settings.as_str(),
279+
output,
280+
)?;
281+
}
282+
},
283+
name if name == execute_name => {
284+
if let Some(overridden_settings) = ExecuteCommand::allowable_field_to_be_overridden(settings) {
285+
queue_permission_override_warning(
286+
allowed_tool.as_str(),
287+
overridden_settings.as_str(),
288+
output,
289+
)?;
290+
}
291+
},
292+
_ => {},
293+
}
294+
}
295+
}
296+
241297
Ok(())
242298
}
243299

@@ -803,6 +859,28 @@ async fn load_legacy_mcp_config(os: &Os) -> eyre::Result<Option<McpServerConfig>
803859
})
804860
}
805861

862+
pub fn queue_permission_override_warning(
863+
tool_name: &str,
864+
overridden_settings: &str,
865+
output: &mut impl Write,
866+
) -> Result<(), std::io::Error> {
867+
Ok(queue!(
868+
output,
869+
style::SetForegroundColor(Color::Yellow),
870+
style::Print("WARN: "),
871+
style::ResetColor,
872+
style::Print("You have trusted "),
873+
style::SetForegroundColor(Color::Green),
874+
style::Print(tool_name),
875+
style::ResetColor,
876+
style::Print(" tool, which overrides the toolsSettings: "),
877+
style::SetForegroundColor(Color::Cyan),
878+
style::Print(overridden_settings),
879+
style::ResetColor,
880+
style::Print("\n"),
881+
)?)
882+
}
883+
806884
fn default_schema() -> String {
807885
"https://raw.githubusercontent.com/aws/amazon-q-developer-cli/refs/heads/main/schemas/agent-v1.json".into()
808886
}

‎crates/chat-cli/src/cli/chat/mod.rs‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,10 @@ use util::{
123123
use winnow::Partial;
124124
use winnow::stream::Offset;
125125

126-
use super::agent::PermissionEvalResult;
126+
use super::agent::{
127+
DEFAULT_AGENT_NAME,
128+
PermissionEvalResult,
129+
};
127130
use crate::api_client::model::ToolResultStatus;
128131
use crate::api_client::{
129132
self,
@@ -611,7 +614,7 @@ impl ChatSession {
611614
": cannot resume conversation with {profile} because it no longer exists. Using default.\n"
612615
))
613616
)?;
614-
let _ = agents.switch("default");
617+
let _ = agents.switch(DEFAULT_AGENT_NAME);
615618
}
616619
}
617620
cs.agents = agents;
@@ -622,6 +625,10 @@ impl ChatSession {
622625
false => ConversationState::new(conversation_id, agents, tool_config, tool_manager, model_id, os).await,
623626
};
624627

628+
if let Some(agent) = conversation.agents.get_active() {
629+
agent.validate_tool_settings(&mut stderr)?;
630+
}
631+
625632
// Spawn a task for listening and broadcasting sigints.
626633
let (ctrlc_tx, ctrlc_rx) = tokio::sync::broadcast::channel(4);
627634
tokio::spawn(async move {
@@ -1739,6 +1746,12 @@ impl ChatSession {
17391746
.clone()
17401747
.unwrap_or(tool_use.name.clone());
17411748
self.conversation.agents.trust_tools(vec![formatted_tool_name]);
1749+
1750+
if let Some(agent) = self.conversation.agents.get_active() {
1751+
agent
1752+
.validate_tool_settings(&mut self.stderr)
1753+
.map_err(|_e| ChatError::Custom("Failed to validate agent tool settings".into()))?;
1754+
}
17421755
}
17431756
tool_use.accepted = true;
17441757

‎crates/chat-cli/src/cli/chat/tools/execute/mod.rs‎

Lines changed: 26 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,12 @@ impl ExecuteCommand {
176176
Ok(())
177177
}
178178

179+
pub fn allowable_field_to_be_overridden(settings: &serde_json::Value) -> Option<String> {
180+
settings
181+
.get("allowedCommands")
182+
.map(|value| format!("allowedCommands: {}", value))
183+
}
184+
179185
pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult {
180186
#[derive(Debug, Deserialize)]
181187
#[serde(rename_all = "camelCase")]
@@ -196,7 +202,7 @@ impl ExecuteCommand {
196202
let tool_name = if cfg!(windows) { "execute_cmd" } else { "execute_bash" };
197203
let is_in_allowlist = agent.allowed_tools.contains(tool_name);
198204
match agent.tools_settings.get(tool_name) {
199-
Some(settings) if is_in_allowlist => {
205+
Some(settings) => {
200206
let Settings {
201207
allowed_commands,
202208
denied_commands,
@@ -220,7 +226,7 @@ impl ExecuteCommand {
220226
return PermissionEvalResult::Deny(denied_match_set);
221227
}
222228

223-
if self.requires_acceptance(Some(&allowed_commands), allow_read_only) {
229+
if !is_in_allowlist || self.requires_acceptance(Some(&allowed_commands), allow_read_only) {
224230
PermissionEvalResult::Ask
225231
} else {
226232
PermissionEvalResult::Allow
@@ -257,10 +263,7 @@ pub fn format_output(output: &str, max_size: usize) -> String {
257263

258264
#[cfg(test)]
259265
mod tests {
260-
use std::collections::{
261-
HashMap,
262-
HashSet,
263-
};
266+
use std::collections::HashMap;
264267

265268
use super::*;
266269
use crate::cli::agent::ToolSettingTarget;
@@ -402,13 +405,8 @@ mod tests {
402405
#[test]
403406
fn test_eval_perm() {
404407
let tool_name = if cfg!(windows) { "execute_cmd" } else { "execute_bash" };
405-
let agent = Agent {
408+
let mut agent = Agent {
406409
name: "test_agent".to_string(),
407-
allowed_tools: {
408-
let mut allowed_tools = HashSet::<String>::new();
409-
allowed_tools.insert(tool_name.to_string());
410-
allowed_tools
411-
},
412410
tools_settings: {
413411
let mut map = HashMap::<ToolSettingTarget, serde_json::Value>::new();
414412
map.insert(
@@ -422,20 +420,32 @@ mod tests {
422420
..Default::default()
423421
};
424422

425-
let tool = serde_json::from_value::<ExecuteCommand>(serde_json::json!({
423+
let tool_one = serde_json::from_value::<ExecuteCommand>(serde_json::json!({
426424
"command": "git status",
427425
}))
428426
.unwrap();
429427

430-
let res = tool.eval_perm(&agent);
428+
let res = tool_one.eval_perm(&agent);
431429
assert!(matches!(res, PermissionEvalResult::Deny(ref rules) if rules.contains(&"\\Agit .*\\z".to_string())));
432430

433-
let tool = serde_json::from_value::<ExecuteCommand>(serde_json::json!({
431+
let tool_two = serde_json::from_value::<ExecuteCommand>(serde_json::json!({
434432
"command": "echo hello",
435433
}))
436434
.unwrap();
437435

438-
let res = tool.eval_perm(&agent);
436+
let res = tool_two.eval_perm(&agent);
437+
assert!(matches!(res, PermissionEvalResult::Ask));
438+
439+
agent.allowed_tools.insert(tool_name.to_string());
440+
441+
let res = tool_two.eval_perm(&agent);
442+
assert!(matches!(res, PermissionEvalResult::Allow));
443+
444+
// Denied list should remain denied
445+
let res = tool_one.eval_perm(&agent);
446+
assert!(matches!(res, PermissionEvalResult::Deny(ref rules) if rules.contains(&"\\Agit .*\\z".to_string())));
447+
448+
let res = tool_two.eval_perm(&agent);
439449
assert!(matches!(res, PermissionEvalResult::Allow));
440450
}
441451

‎crates/chat-cli/src/cli/chat/tools/fs_read.rs‎

Lines changed: 26 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,12 @@ impl FsRead {
102102
}
103103
}
104104

105+
pub fn allowable_field_to_be_overridden(settings: &serde_json::Value) -> Option<String> {
106+
settings
107+
.get("allowedPaths")
108+
.map(|value| format!("allowedPaths: {}", value))
109+
}
110+
105111
pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult {
106112
#[derive(Debug, Deserialize)]
107113
#[serde(rename_all = "camelCase")]
@@ -120,7 +126,7 @@ impl FsRead {
120126

121127
let is_in_allowlist = agent.allowed_tools.contains("fs_read");
122128
match agent.tools_settings.get("fs_read") {
123-
Some(settings) if is_in_allowlist => {
129+
Some(settings) => {
124130
let Settings {
125131
allowed_paths,
126132
denied_paths,
@@ -182,7 +188,7 @@ impl FsRead {
182188

183189
// We only want to ask if we are not allowing read only
184190
// operation
185-
if !allow_read_only && !allow_set.is_match(path) {
191+
if !is_in_allowlist && !allow_read_only && !allow_set.is_match(path) {
186192
ask = true;
187193
}
188194
},
@@ -203,7 +209,10 @@ impl FsRead {
203209

204210
// We only want to ask if we are not allowing read only
205211
// operation
206-
if !allow_read_only && !paths.iter().any(|path| allow_set.is_match(path)) {
212+
if !is_in_allowlist
213+
&& !allow_read_only
214+
&& !paths.iter().any(|path| allow_set.is_match(path))
215+
{
207216
ask = true;
208217
}
209218
},
@@ -838,10 +847,7 @@ fn format_mode(mode: u32) -> [char; 9] {
838847

839848
#[cfg(test)]
840849
mod tests {
841-
use std::collections::{
842-
HashMap,
843-
HashSet,
844-
};
850+
use std::collections::HashMap;
845851

846852
use super::*;
847853
use crate::cli::agent::ToolSettingTarget;
@@ -1381,13 +1387,8 @@ mod tests {
13811387
const DENIED_PATH_ONE: &str = "/some/denied/path";
13821388
const DENIED_PATH_GLOB: &str = "/denied/glob/**/path";
13831389

1384-
let agent = Agent {
1390+
let mut agent = Agent {
13851391
name: "test_agent".to_string(),
1386-
allowed_tools: {
1387-
let mut allowed_tools = HashSet::<String>::new();
1388-
allowed_tools.insert("fs_read".to_string());
1389-
allowed_tools
1390-
},
13911392
tools_settings: {
13921393
let mut map = HashMap::<ToolSettingTarget, serde_json::Value>::new();
13931394
map.insert(
@@ -1401,7 +1402,7 @@ mod tests {
14011402
..Default::default()
14021403
};
14031404

1404-
let tool = serde_json::from_value::<FsRead>(serde_json::json!({
1405+
let tool_one = serde_json::from_value::<FsRead>(serde_json::json!({
14051406
"operations": [
14061407
{ "path": DENIED_PATH_ONE, "mode": "Line", "start_line": 1, "end_line": 2 },
14071408
{ "path": "/denied/glob", "mode": "Directory" },
@@ -1412,7 +1413,17 @@ mod tests {
14121413
}))
14131414
.unwrap();
14141415

1415-
let res = tool.eval_perm(&agent);
1416+
let res = tool_one.eval_perm(&agent);
1417+
assert!(matches!(
1418+
res,
1419+
PermissionEvalResult::Deny(ref deny_list)
1420+
if deny_list.iter().filter(|p| *p == DENIED_PATH_GLOB).collect::<Vec<_>>().len() == 2
1421+
&& deny_list.iter().filter(|p| *p == DENIED_PATH_ONE).collect::<Vec<_>>().len() == 1
1422+
));
1423+
1424+
agent.allowed_tools.insert("fs_read".to_string());
1425+
1426+
let res = tool_one.eval_perm(&agent);
14161427
assert!(matches!(
14171428
res,
14181429
PermissionEvalResult::Deny(ref deny_list)

0 commit comments

Comments
 (0)