Skip to content

Commit 8e3c5b5

Browse files
hayemaxihayemaxi
authored
feat(/tools): update trustall UX, /trust and /untrust capability (#1200)
- Make trustall/acceptall warning more verbose. - Trustall/acceptall warning will show if you start chat with -a or --acceptall or --trust-all-tools - /trust and /untrust takes an unlimited amount of tools at a time - /trust and /untrust will display an error if you specify a tool name that does not exist
1 parent 61c953c commit 8e3c5b5

File tree

2 files changed

+133
-61
lines changed

2 files changed

+133
-61
lines changed

‎crates/q_cli/src/cli/chat/command.rs‎

Lines changed: 29 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use std::collections::HashSet;
12
use std::io::Write;
23

34
use clap::{
@@ -274,8 +275,8 @@ in global or local profiles.
274275

275276
#[derive(Debug, Clone, PartialEq, Eq)]
276277
pub enum ToolsSubcommand {
277-
Trust { tool_name: String },
278-
Untrust { tool_name: String },
278+
Trust { tool_names: HashSet<String> },
279+
Untrust { tool_names: HashSet<String> },
279280
TrustAll,
280281
Reset,
281282
ResetSingle { tool_name: String },
@@ -285,8 +286,8 @@ pub enum ToolsSubcommand {
285286
impl ToolsSubcommand {
286287
const AVAILABLE_COMMANDS: &str = color_print::cstr! {"<cyan!>Available subcommands</cyan!>
287288
<em>help</em> <black!>Show an explanation for the tools command</black!>
288-
<em>trust <<tool name>></em> <black!>Trust a specific tool for the session</black!>
289-
<em>untrust <<tool name>></em> <black!>Revert a tool to per-request confirmation</black!>
289+
<em>trust <<tools...>></em> <black!>Trust a specific tool or tools for the session</black!>
290+
<em>untrust <<tools...>></em> <black!>Revert a tool or tools to per-request confirmation</black!>
290291
<em>trustall</em> <black!>Trust all tools (equivalent to deprecated /acceptall)</black!>
291292
<em>reset</em> <black!>Reset all tools to default permission levels</black!>
292293
<em>reset <<tool name>></em> <black!>Reset a single tool to default permission level</black!>"};
@@ -296,8 +297,8 @@ impl ToolsSubcommand {
296297
Show the current set of tools and their permission setting.
297298
The permission setting states when user confirmation is required. Trusted tools never require confirmation.
298299
Alternatively, specify a subcommand to modify the tool permissions."};
299-
const TRUST_USAGE: &str = "/tools trust <tool name>";
300-
const UNTRUST_USAGE: &str = "/tools untrust <tool name>";
300+
const TRUST_USAGE: &str = "/tools trust <tools...>";
301+
const UNTRUST_USAGE: &str = "/tools untrust <tools...>";
301302

302303
fn usage_msg(header: impl AsRef<str>) -> String {
303304
format!(
@@ -652,25 +653,31 @@ impl Command {
652653

653654
match parts[1].to_lowercase().as_str() {
654655
"trust" => {
655-
let tool_name = parts.get(2);
656-
match tool_name {
657-
Some(tool_name) => Self::Tools {
658-
subcommand: Some(ToolsSubcommand::Trust {
659-
tool_name: (*tool_name).to_string(),
660-
}),
661-
},
662-
None => usage_err!("trust", ToolsSubcommand::TRUST_USAGE),
656+
let mut tool_names = HashSet::new();
657+
for part in &parts[2..] {
658+
tool_names.insert((*part).to_string());
659+
}
660+
661+
if tool_names.is_empty() {
662+
usage_err!("trust", ToolsSubcommand::TRUST_USAGE);
663+
}
664+
665+
Self::Tools {
666+
subcommand: Some(ToolsSubcommand::Trust { tool_names }),
663667
}
664668
},
665669
"untrust" => {
666-
let tool_name = parts.get(2);
667-
match tool_name {
668-
Some(tool_name) => Self::Tools {
669-
subcommand: Some(ToolsSubcommand::Untrust {
670-
tool_name: (*tool_name).to_string(),
671-
}),
672-
},
673-
None => usage_err!("untrust", ToolsSubcommand::UNTRUST_USAGE),
670+
let mut tool_names = HashSet::new();
671+
for part in &parts[2..] {
672+
tool_names.insert((*part).to_string());
673+
}
674+
675+
if tool_names.is_empty() {
676+
usage_err!("untrust", ToolsSubcommand::UNTRUST_USAGE);
677+
}
678+
679+
Self::Tools {
680+
subcommand: Some(ToolsSubcommand::Untrust { tool_names }),
674681
}
675682
},
676683
"trustall" => Self::Tools {

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

Lines changed: 104 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -214,8 +214,8 @@ const HELP_TEXT: &str = color_print::cstr! {"
214214
<em>--summary</em> <black!>Display the summary after compacting</black!>
215215
<em>/tools</em> <black!>View and manage tools and permissions</black!>
216216
<em>help</em> <black!>Show an explanation for the trust command</black!>
217-
<em>trust</em> <black!>Trust a specific tool for the session</black!>
218-
<em>untrust</em> <black!>Revert a tool to per-request confirmation</black!>
217+
<em>trust</em> <black!>Trust a specific tool or tools for the session</black!>
218+
<em>untrust</em> <black!>Revert a tool or tools to per-request confirmation</black!>
219219
<em>trustall</em> <black!>Trust all tools (equivalent to deprecated /acceptall)</black!>
220220
<em>reset</em> <black!>Reset all tools to default permission levels</black!>
221221
<em>/profile</em> <black!>Manage profiles</black!>
@@ -243,6 +243,8 @@ const HELP_TEXT: &str = color_print::cstr! {"
243243
"};
244244

245245
const RESPONSE_TIMEOUT_CONTENT: &str = "Response timed out - message took too long to generate";
246+
const TRUST_ALL_TEXT: &str = color_print::cstr! {"<green!>All tools are now trusted (<red!>!</red!>). Amazon Q will execute tools <bold>without</bold> asking for confirmation.\
247+
\nAgents can sometimes do unexpected things so understand the risks.</green!>"};
246248

247249
pub async fn chat(
248250
input: Option<String>,
@@ -698,6 +700,15 @@ where
698700
)
699701
)?;
700702
execute!(self.output, style::Print("\n"), style::SetForegroundColor(Color::Reset))?;
703+
if self.interactive && self.all_tools_trusted() {
704+
queue!(
705+
self.output,
706+
style::Print(format!(
707+
"{}{TRUST_ALL_TEXT}\n\n",
708+
if !is_small_screen { "\n" } else { "" }
709+
))
710+
)?;
711+
}
701712
self.output.flush()?;
702713

703714
let mut ctrl_c_stream = signal(SignalKind::interrupt())?;
@@ -1934,29 +1945,91 @@ where
19341945
}
19351946
},
19361947
Command::Tools { subcommand } => {
1948+
let existing_tools: HashSet<&String> = self
1949+
.conversation_state
1950+
.tools
1951+
.iter()
1952+
.map(|FigTool::ToolSpecification(spec)| &spec.name)
1953+
.collect();
1954+
19371955
match subcommand {
1938-
Some(ToolsSubcommand::Trust { tool_name }) => {
1939-
self.tool_permissions.trust_tool(&tool_name);
1940-
queue!(
1941-
self.output,
1942-
style::SetForegroundColor(Color::Green),
1943-
style::Print(format!("\nTool '{tool_name}' is now trusted. I will ")),
1944-
style::SetAttribute(Attribute::Bold),
1945-
style::Print("not"),
1946-
style::SetAttribute(Attribute::Reset),
1947-
style::SetForegroundColor(Color::Green),
1948-
style::Print(" ask for confirmation before running this tool."),
1949-
style::SetForegroundColor(Color::Reset),
1950-
)?;
1956+
Some(ToolsSubcommand::Trust { tool_names }) => {
1957+
let (valid_tools, invalid_tools): (Vec<String>, Vec<String>) = tool_names
1958+
.into_iter()
1959+
.partition(|tool_name| existing_tools.contains(tool_name));
1960+
1961+
if !invalid_tools.is_empty() {
1962+
queue!(
1963+
self.output,
1964+
style::SetForegroundColor(Color::Red),
1965+
style::Print(format!("\nCannot trust '{}', ", invalid_tools.join("', '"))),
1966+
if invalid_tools.len() > 1 {
1967+
style::Print("they do not exist.")
1968+
} else {
1969+
style::Print("it does not exist.")
1970+
},
1971+
style::SetForegroundColor(Color::Reset),
1972+
)?;
1973+
}
1974+
if !valid_tools.is_empty() {
1975+
valid_tools.iter().for_each(|t| self.tool_permissions.trust_tool(t));
1976+
queue!(
1977+
self.output,
1978+
style::SetForegroundColor(Color::Green),
1979+
if valid_tools.len() > 1 {
1980+
style::Print(format!("\nTools '{}' are ", valid_tools.join("', '")))
1981+
} else {
1982+
style::Print(format!("\nTool '{}' is ", valid_tools[0]))
1983+
},
1984+
style::Print("now trusted. I will "),
1985+
style::SetAttribute(Attribute::Bold),
1986+
style::Print("not"),
1987+
style::SetAttribute(Attribute::Reset),
1988+
style::SetForegroundColor(Color::Green),
1989+
style::Print(format!(
1990+
" ask for confirmation before running {}.",
1991+
if valid_tools.len() > 1 {
1992+
"these tools"
1993+
} else {
1994+
"this tool"
1995+
}
1996+
)),
1997+
style::SetForegroundColor(Color::Reset),
1998+
)?;
1999+
}
19512000
},
1952-
Some(ToolsSubcommand::Untrust { tool_name }) => {
1953-
self.tool_permissions.untrust_tool(&tool_name);
1954-
queue!(
1955-
self.output,
1956-
style::SetForegroundColor(Color::Green),
1957-
style::Print(format!("\nTool '{tool_name}' set to per-request confirmation."),),
1958-
style::SetForegroundColor(Color::Reset),
1959-
)?;
2001+
Some(ToolsSubcommand::Untrust { tool_names }) => {
2002+
let (valid_tools, invalid_tools): (Vec<String>, Vec<String>) = tool_names
2003+
.into_iter()
2004+
.partition(|tool_name| existing_tools.contains(tool_name));
2005+
2006+
if !invalid_tools.is_empty() {
2007+
queue!(
2008+
self.output,
2009+
style::SetForegroundColor(Color::Red),
2010+
style::Print(format!("\nCannot untrust '{}', ", invalid_tools.join("', '"))),
2011+
if invalid_tools.len() > 1 {
2012+
style::Print("they do not exist.")
2013+
} else {
2014+
style::Print("it does not exist.")
2015+
},
2016+
style::SetForegroundColor(Color::Reset),
2017+
)?;
2018+
}
2019+
if !valid_tools.is_empty() {
2020+
valid_tools.iter().for_each(|t| self.tool_permissions.untrust_tool(t));
2021+
queue!(
2022+
self.output,
2023+
style::SetForegroundColor(Color::Green),
2024+
if valid_tools.len() > 1 {
2025+
style::Print(format!("\nTools '{}' are ", valid_tools.join("', '")))
2026+
} else {
2027+
style::Print(format!("\nTool '{}' is ", valid_tools[0]))
2028+
},
2029+
style::Print("set to per-request confirmation."),
2030+
style::SetForegroundColor(Color::Reset),
2031+
)?;
2032+
}
19602033
},
19612034
Some(ToolsSubcommand::TrustAll) => {
19622035
self.conversation_state
@@ -1965,17 +2038,7 @@ where
19652038
.for_each(|FigTool::ToolSpecification(spec)| {
19662039
self.tool_permissions.trust_tool(spec.name.as_str());
19672040
});
1968-
queue!(
1969-
self.output,
1970-
style::SetForegroundColor(Color::Green),
1971-
style::Print("\nAll tools are now trusted. I will "),
1972-
style::SetAttribute(Attribute::Bold),
1973-
style::Print("not"),
1974-
style::SetAttribute(Attribute::Reset),
1975-
style::SetForegroundColor(Color::Green),
1976-
style::Print(" ask for confirmation before running any tools."),
1977-
style::SetForegroundColor(Color::Reset),
1978-
)?;
2041+
queue!(self.output, style::Print(TRUST_ALL_TEXT),)?;
19792042
},
19802043
Some(ToolsSubcommand::Reset) => {
19812044
self.tool_permissions.reset();
@@ -2864,11 +2927,7 @@ where
28642927

28652928
/// Helper function to generate a prompt based on the current context
28662929
fn generate_tool_trust_prompt(&self) -> String {
2867-
let all_tools_trusted = self.conversation_state.tools.iter().all(|t| match t {
2868-
FigTool::ToolSpecification(t) => self.tool_permissions.is_trusted(&t.name),
2869-
});
2870-
2871-
prompt::generate_prompt(self.conversation_state.current_profile(), all_tools_trusted)
2930+
prompt::generate_prompt(self.conversation_state.current_profile(), self.all_tools_trusted())
28722931
}
28732932

28742933
async fn send_tool_use_telemetry(&mut self) {
@@ -2888,6 +2947,12 @@ where
28882947
(self.terminal_width_provider)().unwrap_or(80)
28892948
}
28902949

2950+
fn all_tools_trusted(&self) -> bool {
2951+
self.conversation_state.tools.iter().all(|t| match t {
2952+
FigTool::ToolSpecification(t) => self.tool_permissions.is_trusted(&t.name),
2953+
})
2954+
}
2955+
28912956
/// Display character limit warnings based on current conversation size
28922957
fn display_char_warnings(&mut self) -> Result<(), std::io::Error> {
28932958
// Check character count and warning level

0 commit comments

Comments
 (0)