diff --git a/crates/chat-cli/src/cli/agent/mod.rs b/crates/chat-cli/src/cli/agent/mod.rs index 3dd297991..793440c38 100644 --- a/crates/chat-cli/src/cli/agent/mod.rs +++ b/crates/chat-cli/src/cli/agent/mod.rs @@ -314,6 +314,7 @@ impl Agent { pub enum PermissionEvalResult { Allow, Ask, + AskWithoutTrust, Deny, } diff --git a/crates/chat-cli/src/cli/chat/mod.rs b/crates/chat-cli/src/cli/chat/mod.rs index 626caa4a6..ec6a96301 100644 --- a/crates/chat-cli/src/cli/chat/mod.rs +++ b/crates/chat-cli/src/cli/chat/mod.rs @@ -516,6 +516,9 @@ pub struct ChatSession { tool_uses: Vec, /// An index into [Self::tool_uses] to represent the current tool use being handled. pending_tool_index: Option, + /// Whether the current pending tool allows trust option (true for Ask, false for + /// AskWithoutTrust) + pending_tool_allows_trust: bool, /// The time immediately after having received valid tool uses from the model. /// /// Used to track the time taken from initially prompting the user to tool execute @@ -646,6 +649,7 @@ impl ChatSession { tool_uses: vec![], user_turn_request_metadata: vec![], pending_tool_index: None, + pending_tool_allows_trust: false, tool_turn_start_time: None, tool_use_telemetry_events: HashMap::new(), tool_use_status: ToolUseStatus::Idle, @@ -1003,6 +1007,7 @@ impl ChatSession { self.conversation.enforce_conversation_invariants(); self.conversation.reset_next_user_message(); self.pending_tool_index = None; + self.pending_tool_allows_trust = false; self.tool_turn_start_time = None; self.reset_user_turn(); @@ -1498,28 +1503,47 @@ impl ChatSession { let show_tool_use_confirmation_dialog = !skip_printing_tools && self.pending_tool_index.is_some(); if show_tool_use_confirmation_dialog { - execute!( - self.stderr, - style::SetForegroundColor(Color::DarkGrey), - style::Print("\nAllow this action? Use '"), - style::SetForegroundColor(Color::Green), - style::Print("t"), - style::SetForegroundColor(Color::DarkGrey), - style::Print("' to trust (always allow) this tool for the session. ["), - style::SetForegroundColor(Color::Green), - style::Print("y"), - style::SetForegroundColor(Color::DarkGrey), - style::Print("/"), - style::SetForegroundColor(Color::Green), - style::Print("n"), - style::SetForegroundColor(Color::DarkGrey), - style::Print("/"), - style::SetForegroundColor(Color::Green), - style::Print("t"), - style::SetForegroundColor(Color::DarkGrey), - style::Print("]:\n\n"), - style::SetForegroundColor(Color::Reset), - )?; + if self.pending_tool_allows_trust { + // Show full prompt with trust option for Ask + execute!( + self.stderr, + style::SetForegroundColor(Color::DarkGrey), + style::Print("\nAllow this action? Use '"), + style::SetForegroundColor(Color::Green), + style::Print("t"), + style::SetForegroundColor(Color::DarkGrey), + style::Print("' to trust (always allow) this tool for the session. ["), + style::SetForegroundColor(Color::Green), + style::Print("y"), + style::SetForegroundColor(Color::DarkGrey), + style::Print("/"), + style::SetForegroundColor(Color::Green), + style::Print("n"), + style::SetForegroundColor(Color::DarkGrey), + style::Print("/"), + style::SetForegroundColor(Color::Green), + style::Print("t"), + style::SetForegroundColor(Color::DarkGrey), + style::Print("]:\n\n"), + style::SetForegroundColor(Color::Reset), + )?; + } else { + // Show simplified prompt without trust option for AskWithoutTrust + execute!( + self.stderr, + style::SetForegroundColor(Color::DarkGrey), + style::Print("\nAllow this action? ["), + style::SetForegroundColor(Color::Green), + style::Print("y"), + style::SetForegroundColor(Color::DarkGrey), + style::Print("/"), + style::SetForegroundColor(Color::Green), + style::Print("n"), + style::SetForegroundColor(Color::DarkGrey), + style::Print("]:\n\n"), + style::SetForegroundColor(Color::Reset), + )?; + } } // Do this here so that the skim integration sees an updated view of the context *during the current @@ -1712,7 +1736,7 @@ impl ChatSession { } else { // Check for a pending tool approval if let Some(index) = self.pending_tool_index { - let is_trust = ["t", "T"].contains(&input); + let is_trust = ["t", "T"].contains(&input) && self.pending_tool_allows_trust; let tool_use = &mut self.tool_uses[index]; if ["y", "Y"].contains(&input) || is_trust { if is_trust { @@ -1793,13 +1817,18 @@ impl ChatSession { } let mut denied = false; + let mut allows_trust = false; let allowed = self.conversation .agents .get_active() .is_some_and(|a| match tool.tool.requires_acceptance(a) { PermissionEvalResult::Allow => true, - PermissionEvalResult::Ask => false, + PermissionEvalResult::Ask => { + allows_trust = true; + false + }, + PermissionEvalResult::AskWithoutTrust => false, PermissionEvalResult::Deny => { denied = true; false @@ -1839,6 +1868,7 @@ impl ChatSession { } self.pending_tool_index = Some(i); + self.pending_tool_allows_trust = allows_trust; return Ok(ChatState::PromptUser { skip_printing_tools: false, @@ -2275,6 +2305,7 @@ impl ChatSession { } else { self.tool_uses.clear(); self.pending_tool_index = None; + self.pending_tool_allows_trust = false; self.tool_turn_start_time = None; self.send_chat_telemetry(os, TelemetryResult::Succeeded, None, None, None, true) @@ -2398,6 +2429,7 @@ impl ChatSession { self.conversation.enforce_conversation_invariants(); self.conversation.reset_next_user_message(); self.pending_tool_index = None; + self.pending_tool_allows_trust = false; self.tool_turn_start_time = None; return Ok(ChatState::PromptUser { skip_printing_tools: false, diff --git a/crates/chat-cli/src/cli/chat/tools/use_aws.rs b/crates/chat-cli/src/cli/chat/tools/use_aws.rs index 59b41e8b0..2f0d6d3e2 100644 --- a/crates/chat-cli/src/cli/chat/tools/use_aws.rs +++ b/crates/chat-cli/src/cli/chat/tools/use_aws.rs @@ -221,10 +221,9 @@ impl UseAws { } PermissionEvalResult::Ask }, - None if is_in_allowlist => PermissionEvalResult::Allow, _ => { if self.requires_acceptance() { - PermissionEvalResult::Ask + PermissionEvalResult::AskWithoutTrust } else { PermissionEvalResult::Allow } @@ -324,6 +323,63 @@ mod tests { ); } + #[test] + fn test_trusted_tool_still_requires_acceptance_for_write_ops() { + use crate::cli::agent::Agent; + + // Create an agent with use_aws in the allowed tools (trusted) + let mut agent = Agent::default(); + agent.allowed_tools.insert("use_aws".to_string()); + + // Test a write operation (should still require acceptance even if trusted) + let write_cmd = use_aws! {{ + "service_name": "s3", + "operation_name": "put-object", + "region": "us-west-2", + "profile_name": "default", + "label": "" + }}; + assert_eq!(write_cmd.eval_perm(&agent), PermissionEvalResult::AskWithoutTrust); + + // Test a read operation (should be allowed for trusted tools) + let read_cmd = use_aws! {{ + "service_name": "s3", + "operation_name": "list-objects", + "region": "us-west-2", + "profile_name": "default", + "label": "" + }}; + assert_eq!(read_cmd.eval_perm(&agent), PermissionEvalResult::Allow); + } + + #[test] + fn test_untrusted_tool_uses_ask_without_trust_for_write_ops() { + use crate::cli::agent::Agent; + + // Create an agent without use_aws in the allowed tools (untrusted) + let agent = Agent::default(); + + // Test a write operation (should use AskWithoutTrust for untrusted tools) + let write_cmd = use_aws! {{ + "service_name": "s3", + "operation_name": "put-object", + "region": "us-west-2", + "profile_name": "default", + "label": "" + }}; + assert_eq!(write_cmd.eval_perm(&agent), PermissionEvalResult::AskWithoutTrust); + + // Test a read operation (should be allowed for untrusted tools too) + let read_cmd = use_aws! {{ + "service_name": "s3", + "operation_name": "list-objects", + "region": "us-west-2", + "profile_name": "default", + "label": "" + }}; + assert_eq!(read_cmd.eval_perm(&agent), PermissionEvalResult::Allow); + } + #[tokio::test] #[ignore = "not in ci"] async fn test_aws_output() {