feat: implement database setting for trusting tools

Author: fsmiamoto

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

@@ -19,121 +19,52 @@ pub mod tools;
 pub mod util;
 
 use std::borrow::Cow;
-use std::collections::{
-    HashMap,
-    HashSet,
-    VecDeque,
-};
+use std::collections::{HashMap, HashSet, VecDeque};
 use std::io::Write;
 use std::process::ExitCode;
 use std::time::Duration;
 
 use amzn_codewhisperer_client::types::SubscriptionStatus;
-use clap::{
-    Args,
-    CommandFactory,
-    Parser,
-};
+use clap::{Args, CommandFactory, Parser};
 use context::ContextManager;
 pub use conversation::ConversationState;
 use conversation::TokenWarningLevel;
-use crossterm::style::{
-    Attribute,
-    Color,
-    Stylize,
-};
-use crossterm::{
-    cursor,
-    execute,
-    queue,
-    style,
-    terminal,
-};
-use eyre::{
-    Report,
-    Result,
-    bail,
-    eyre,
-};
+use crossterm::style::{Attribute, Color, Stylize};
+use crossterm::{cursor, execute, queue, style, terminal};
+use eyre::{Report, Result, bail, eyre};
 use input_source::InputSource;
-use message::{
-    AssistantMessage,
-    AssistantToolUse,
-    ToolUseResult,
-    ToolUseResultBlock,
-};
-use parse::{
-    ParseState,
-    interpret_markdown,
-};
-use parser::{
-    RecvErrorKind,
-    ResponseParser,
-};
+use message::{AssistantMessage, AssistantToolUse, ToolUseResult, ToolUseResultBlock};
+use parse::{ParseState, interpret_markdown};
+use parser::{RecvErrorKind, ResponseParser};
 use regex::Regex;
-use spinners::{
-    Spinner,
-    Spinners,
-};
+use spinners::{Spinner, Spinners};
 use thiserror::Error;
 use time::OffsetDateTime;
 use token_counter::TokenCounter;
 use tokio::signal::ctrl_c;
-use tool_manager::{
-    McpServerConfig,
-    ToolManager,
-    ToolManagerBuilder,
-};
+use tool_manager::{McpServerConfig, ToolManager, ToolManagerBuilder};
 use tools::gh_issue::GhIssueContext;
-use tools::{
-    OutputKind,
-    QueuedTool,
-    Tool,
-    ToolPermissions,
-    ToolSpec,
-};
-use tracing::{
-    debug,
-    error,
-    info,
-    trace,
-    warn,
-};
+use tools::{OutputKind, QueuedTool, Tool, ToolPermissions, ToolSpec};
+use tracing::{debug, error, info, trace, warn};
 use util::images::RichImageBlock;
 use util::ui::draw_box;
-use util::{
-    animate_output,
-    play_notification_bell,
-};
+use util::{animate_output, play_notification_bell};
 use winnow::Partial;
 use winnow::stream::Offset;
 
 use crate::api_client::ApiClientError;
-use crate::api_client::model::{
-    Tool as FigTool,
-    ToolResultStatus,
-};
+use crate::api_client::model::{Tool as FigTool, ToolResultStatus};
 use crate::api_client::send_message_output::SendMessageOutput;
 use crate::auth::AuthError;
 use crate::auth::builder_id::is_idc_user;
 use crate::cli::chat::cli::SlashCommand;
-use crate::cli::chat::cli::model::{
-    MODEL_OPTIONS,
-    default_model_id,
-};
-use crate::cli::chat::cli::prompts::{
-    GetPromptError,
-    PromptsSubcommand,
-};
+use crate::cli::chat::cli::model::{MODEL_OPTIONS, default_model_id};
+use crate::cli::chat::cli::prompts::{GetPromptError, PromptsSubcommand};
 use crate::database::settings::Setting;
 use crate::mcp_client::Prompt;
 use crate::os::Os;
 use crate::telemetry::core::ToolUseEventBuilder;
-use crate::telemetry::{
-    ReasonCode,
-    TelemetryResult,
-    get_error_reason,
-};
+use crate::telemetry::{ReasonCode, TelemetryResult, get_error_reason};
 
 const LIMIT_REACHED_TEXT: &str = color_print::cstr! { "You've used all your free requests for this month. You have two options:
 1. Upgrade to a paid subscription for increased limits. See our Pricing page for what's included> <blue!>https://aws.amazon.com/q/developer/pricing/</blue!>
@@ -255,6 +186,7 @@ impl ChatArgs {
             .build(os, Box::new(std::io::stderr()), !self.non_interactive)
             .await?;
         let tool_config = tool_manager.load_tools(os, &mut stderr).await?;
+
         let mut tool_permissions = ToolPermissions::new(tool_config.len());
 
         if self.trust_all_tools {
@@ -279,6 +211,9 @@ impl ChatArgs {
                     tool_permissions.untrust_tool(&tool.name);
                 }
             }
+        } else {
+            // CLI args has precendence over Database config
+            tool_permissions.trust_from_database(&os.database);
         }
 
         ChatSession::new(

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

@@ -7,38 +7,27 @@ pub mod knowledge;
 pub mod thinking;
 pub mod use_aws;
 
-use std::collections::{
-    HashMap,
-    HashSet,
-};
+use std::collections::{HashMap, HashSet};
 use std::io::Write;
-use std::path::{
-    Path,
-    PathBuf,
-};
+use std::path::{Path, PathBuf};
 
 use crossterm::queue;
-use crossterm::style::{
-    self,
-    Color,
-    Stylize,
-};
+use crossterm::style::{self, Color, Stylize};
 use custom_tool::CustomTool;
 use execute::ExecuteCommand;
 use eyre::Result;
 use fs_read::FsRead;
 use fs_write::FsWrite;
 use gh_issue::GhIssue;
 use knowledge::Knowledge;
-use serde::{
-    Deserialize,
-    Serialize,
-};
+use serde::{Deserialize, Serialize};
 use thinking::Thinking;
 use use_aws::UseAws;
 
 use super::consts::MAX_TOOL_RESPONSE_SIZE;
 use super::util::images::RichImageBlocks;
+use crate::database::Database;
+use crate::database::settings::Setting;
 use crate::os::Os;
 
 /// Represents an executable tool use.
@@ -224,6 +213,15 @@ impl ToolPermissions {
         self.permissions.contains_key(tool_name)
     }
 
+    pub fn trust_from_database(&mut self, database: &Database) {
+        database
+            .settings
+            .get_array::<String>(Setting::TrustedTools)
+            .into_iter()
+            .flatten()
+            .for_each(|tool_name| self.trust_tool(&tool_name));
+    }
+
     /// Provide default permission labels for the built-in set of tools.
     // This "static" way avoids needing to construct a tool instance.
     fn default_permission_label(&self, tool_name: &str) -> String {

crates/chat-cli/src/database/settings.rs

@@ -2,16 +2,9 @@ use std::fmt::Display;
 use std::io::SeekFrom;
 
 use fd_lock::RwLock;
-use serde_json::{
-    Map,
-    Value,
-};
+use serde_json::{Map, Value};
 use tokio::fs::File;
-use tokio::io::{
-    AsyncReadExt,
-    AsyncSeekExt,
-    AsyncWriteExt,
-};
+use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
 
 use super::DatabaseError;
 
@@ -33,6 +26,7 @@ pub enum Setting {
     McpNoInteractiveTimeout,
     McpLoadedBefore,
     ChatDefaultModel,
+    TrustedTools,
 }
 
 impl AsRef<str> for Setting {
@@ -54,6 +48,7 @@ impl AsRef<str> for Setting {
             Self::McpNoInteractiveTimeout => "mcp.noInteractiveTimeout",
             Self::McpLoadedBefore => "mcp.loadedBefore",
             Self::ChatDefaultModel => "chat.defaultModel",
+            Self::TrustedTools => "tools.trusted",
         }
     }
 }
@@ -85,6 +80,7 @@ impl TryFrom<&str> for Setting {
             "mcp.noInteractiveTimeout" => Ok(Self::McpNoInteractiveTimeout),
             "mcp.loadedBefore" => Ok(Self::McpLoadedBefore),
             "chat.defaultModel" => Ok(Self::ChatDefaultModel),
+            "tools.trusted" => Ok(Self::TrustedTools),
             _ => Err(DatabaseError::InvalidSetting(value.to_string())),
         }
     }
@@ -150,6 +146,19 @@ impl Settings {
         self.get(key).and_then(|value| value.as_str().map(|s| s.into()))
     }
 
+    pub fn get_array<T>(&self, key: Setting) -> Option<Vec<T>>
+    where
+        T: serde::de::DeserializeOwned,
+    {
+        self.get(key).and_then(|value| {
+            value.as_array().and_then(|arr| {
+                arr.iter()
+                    .map(|v| serde_json::from_value(v.clone()).ok())
+                    .collect::<Option<Vec<T>>>()
+            })
+        })
+    }
+
     pub fn get_int(&self, key: Setting) -> Option<i64> {
         self.get(key).and_then(|value| value.as_i64())
     }
@@ -204,12 +213,17 @@ mod test {
         assert_eq!(settings.get(Setting::ShareCodeWhispererContent), None);
         assert_eq!(settings.get(Setting::McpLoadedBefore), None);
         assert_eq!(settings.get(Setting::ChatDefaultModel), None);
+        assert_eq!(settings.get(Setting::TrustedTools), None);
 
         settings.set(Setting::TelemetryEnabled, true).await.unwrap();
         settings.set(Setting::OldClientId, "test").await.unwrap();
         settings.set(Setting::ShareCodeWhispererContent, false).await.unwrap();
         settings.set(Setting::McpLoadedBefore, true).await.unwrap();
         settings.set(Setting::ChatDefaultModel, "model 1").await.unwrap();
+        settings
+            .set(Setting::TrustedTools, r#"["tool_a","tool_b"]"#)
+            .await
+            .unwrap();
 
         assert_eq!(settings.get(Setting::TelemetryEnabled), Some(&Value::Bool(true)));
         assert_eq!(
@@ -225,15 +239,21 @@ mod test {
             settings.get(Setting::ChatDefaultModel),
             Some(&Value::String("model 1".to_string()))
         );
+        assert_eq!(
+            settings.get(Setting::TrustedTools),
+            Some(&Value::String(r#"["tool_a","tool_b"]"#.to_string()))
+        );
 
         settings.remove(Setting::TelemetryEnabled).await.unwrap();
         settings.remove(Setting::OldClientId).await.unwrap();
         settings.remove(Setting::ShareCodeWhispererContent).await.unwrap();
         settings.remove(Setting::McpLoadedBefore).await.unwrap();
+        settings.remove(Setting::TrustedTools).await.unwrap();
 
         assert_eq!(settings.get(Setting::TelemetryEnabled), None);
         assert_eq!(settings.get(Setting::OldClientId), None);
         assert_eq!(settings.get(Setting::ShareCodeWhispererContent), None);
         assert_eq!(settings.get(Setting::McpLoadedBefore), None);
+        assert_eq!(settings.get(Setting::TrustedTools), None);
     }
 }