fix: aws parameter bugs (#688)

Author: brandonskiser

Cargo.lock

@@ -32,7 +32,9 @@ anstream = "0.6.13"
 anyhow = "1.0.75"
 appkit-nsworkspace-bindings = { path = "crates/macos-utils/appkit-nsworkspace-bindings" }
 async-trait = "0.1.74"
-aws-types = "1.2.0"
+aws-smithy-runtime-api = "1.6.1"
+aws-smithy-types = "1.2.10"
+aws-types = "1.3.0"
 base64 = "0.22.1"
 block2 = "0.5.1"
 bytes = "1.5.0"
@@ -48,6 +50,7 @@ clap = { version = "4.5.23", features = [
     "wrap_help",
 ] }
 cocoa = "0.26.0"
+convert_case = "0.6.0"
 core-foundation = "0.10.0"
 core-foundation-sys = "0.8.7"
 core-graphics = "0.24.0"

Cargo.toml

@@ -22,9 +22,9 @@ amzn-qdeveloper-streaming-client.workspace = true
 aws-config = "1.0.3"
 aws-credential-types = "1.0.3"
 aws-smithy-async = "1.2.2"
-aws-smithy-runtime-api = "1.6.1"
-aws-smithy-types = "1.2.10"
-aws-types = "1.3.0"
+aws-smithy-runtime-api.workspace = true
+aws-smithy-types.workspace = true
+aws-types.workspace = true
 bytes.workspace = true
 fig_auth.workspace = true
 fig_aws_common.workspace = true

crates/fig_api_client/Cargo.toml

@@ -9,10 +9,10 @@ license.workspace = true
 
 [dependencies]
 aws-runtime = "1.4.4"
-aws-smithy-runtime-api = "1.6.1"
-aws-smithy-types = "1.2.10"
-aws-types = "1.3.0"
-http = "0.2.12"
+aws-smithy-runtime-api.workspace = true
+aws-smithy-types.workspace = true
+aws-types.workspace = true
+http.workspace = true
 tracing.workspace = true
 
 [dev-dependencies]

crates/fig_aws_common/Cargo.toml

@@ -126,7 +126,7 @@ impl Intercept for UserAgentOverrideInterceptor {
         )?;
 
         let headers = context.request_mut().headers_mut();
-        headers.insert(USER_AGENT, ua.aws_ua_header());
+        headers.insert(USER_AGENT.as_str(), ua.aws_ua_header());
         Ok(())
     }
 }

crates/fig_aws_common/src/user_agent_override_interceptor.rs

@@ -34,6 +34,7 @@ clap_complete = "4.5.39"
 clap_complete_fig = "4.4.0"
 color-eyre = "0.6.2"
 color-print = "0.3.5"
+convert_case.workspace = true
 crossterm = { version = "0.28.1", features = ["event-stream"] }
 ctrlc = "3.2.5"
 dialoguer = { version = "0.11.0", features = ["fuzzy-select"] }

crates/q_cli/Cargo.toml

@@ -75,7 +75,7 @@
   },
   {
     "name": "use_aws",
-    "description": "Make an AWS CLI api call with the specified service, operation, and parameters. The arguments MUST conform to the AWS CLI specification.",
+    "description": "Make an AWS CLI api call with the specified service, operation, and parameters. All arguments MUST conform to the AWS CLI specification.",
     "input_schema": {
       "type": "object",
       "properties": {
@@ -89,7 +89,7 @@
         },
         "parameters": {
           "type": "object",
-          "description": "The parameters for the operation. Each parameter must be encoded as a string."
+          "description": "The parameters for the operation. The parameter keys MUST conform to the AWS CLI specification. You should prefer to use JSON Syntax over shorthand syntax wherever possible."
         },
         "region": {
           "type": "string",

crates/q_cli/src/cli/chat/tools/tool_index.json

@@ -3,6 +3,10 @@ use std::io::Write;
 use std::process::Stdio;
 
 use bstr::ByteSlice;
+use convert_case::{
+    Case,
+    Casing,
+};
 use crossterm::{
     queue,
     style,
@@ -19,23 +23,23 @@ use super::{
     OutputKind,
 };
 
-const ALLOWED_OPS: [&str; 6] = ["get", "describe", "list", "ls", "search", "batch_get"];
+const READONLY_OPS: [&str; 6] = ["get", "describe", "list", "ls", "search", "batch_get"];
 
 // TODO: we should perhaps composite this struct with an interface that we can use to mock the
 // actual cli with. That will allow us to more thoroughly test it.
 #[derive(Debug, Deserialize)]
 pub struct UseAws {
     pub service_name: String,
     pub operation_name: String,
-    pub parameters: Option<HashMap<String, String>>,
+    pub parameters: Option<HashMap<String, serde_json::Value>>,
     pub region: String,
     pub profile_name: Option<String>,
     pub label: Option<String>,
 }
 
 impl UseAws {
     pub fn requires_consent(&self) -> bool {
-        !ALLOWED_OPS.iter().any(|op| self.operation_name.starts_with(op))
+        !READONLY_OPS.iter().any(|op| self.operation_name.starts_with(op))
     }
 
     pub async fn invoke(&self, _ctx: &Context, _updates: impl Write) -> Result<InvokeOutput> {
@@ -45,15 +49,9 @@ impl UseAws {
             command.arg("--profile").arg(profile_name);
         }
         command.arg(&self.service_name).arg(&self.operation_name);
-        if let Some(parameters) = &self.parameters {
-            for (param_name, val) in parameters {
-                // Model might sometimes give parameters capitalized.
-                let param_name = param_name.to_lowercase();
-                if param_name.starts_with("--") {
-                    command.arg(param_name).arg(val);
-                } else {
-                    command.arg(format!("--{}", param_name)).arg(val);
-                }
+        if let Some(parameters) = self.cli_parameters() {
+            for (name, val) in parameters {
+                command.arg(name).arg(val);
             }
         }
         let output = command
@@ -87,11 +85,11 @@ impl UseAws {
             style::Print("Running aws cli command:\n"),
             style::Print(format!("Service name: {}\n", self.service_name)),
             style::Print(format!("Operation name: {}\n", self.operation_name)),
-            style::Print("Parameters: \n".to_string()),
         )?;
         if let Some(parameters) = &self.parameters {
+            queue!(updates, style::Print("Parameters: \n".to_string()))?;
             for (name, value) in parameters {
-                queue!(updates, style::Print(format!("{}: {}\n", name, value)))?;
+                queue!(updates, style::Print(format!("- {}: {}\n", name, value)))?;
             }
         }
 
@@ -112,6 +110,22 @@ impl UseAws {
     pub async fn validate(&mut self, _ctx: &Context) -> Result<()> {
         Ok(())
     }
+
+    /// Returns the CLI arguments properly formatted as kebab case if parameters is
+    /// [Option::Some], otherwise None
+    fn cli_parameters(&self) -> Option<Vec<(String, String)>> {
+        if let Some(parameters) = &self.parameters {
+            let mut params = vec![];
+            for (param_name, val) in parameters {
+                let param_name = format!("--{}", param_name.trim_start_matches("--").to_case(Case::Kebab));
+                let param_val = val.as_str().map(|s| s.to_string()).unwrap_or(val.to_string());
+                params.push((param_name, param_val));
+            }
+            Some(params)
+        } else {
+            None
+        }
+    }
 }
 
 #[cfg(test)]
@@ -152,6 +166,34 @@ mod tests {
         assert!(cmd.requires_consent());
     }
 
+    #[test]
+    fn test_use_aws_deser() {
+        let cmd = use_aws! {{
+            "service_name": "s3",
+            "operation_name": "put-object",
+            "parameters": {
+                "TableName": "table-name",
+                "KeyConditionExpression": "PartitionKey = :pkValue"
+            },
+            "region": "us-west-2",
+            "profile_name": "default",
+            "label": ""
+        }};
+        let params = cmd.cli_parameters().unwrap();
+        assert!(
+            params.iter().any(|p| p.0 == "--table-name" && p.1 == "table-name"),
+            "not found in {:?}",
+            params
+        );
+        assert!(
+            params
+                .iter()
+                .any(|p| p.0 == "--key-condition-expression" && p.1 == "PartitionKey = :pkValue"),
+            "not found in {:?}",
+            params
+        );
+    }
+
     #[tokio::test]
     #[ignore = "not in ci"]
     async fn test_aws_read_only() {