feat(context): context hooks with /context hooks

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

@@ -1,5 +1,9 @@
 use std::io::Write;
 
+use clap::{
+    Parser,
+    Subcommand,
+};
 use crossterm::style::Color;
 use crossterm::{
     queue,
@@ -87,6 +91,57 @@ Profiles allow you to organize and manage different sets of context files for di
     }
 }
 
+#[derive(Parser, Debug, Clone)]
+#[command(name = "hooks", disable_help_flag = true, disable_help_subcommand = true)]
+struct HooksCommand {
+    #[command(subcommand)]
+    command: HooksSubcommand,
+}
+
+#[derive(Subcommand, Debug, Clone, Eq, PartialEq)]
+pub enum HooksSubcommand {
+    Add {
+        name: String,
+
+        #[arg(long, value_parser = ["per_prompt", "conversation_start"])]
+        r#type: String,
+
+        #[arg(long, value_parser = clap::value_parser!(String))]
+        command: String,
+
+        #[arg(long)]
+        global: bool,
+    },
+    #[command(name = "rm")]
+    Remove {
+        name: String,
+
+        #[arg(long)]
+        global: bool,
+    },
+    Enable {
+        name: String,
+
+        #[arg(long)]
+        global: bool,
+    },
+    Disable {
+        name: String,
+
+        #[arg(long)]
+        global: bool,
+    },
+    EnableAll {
+        #[arg(long)]
+        global: bool,
+    },
+    DisableAll {
+        #[arg(long)]
+        global: bool,
+    },
+    Help,
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum ContextSubcommand {
     Show {
@@ -104,6 +159,9 @@ pub enum ContextSubcommand {
     Clear {
         global: bool,
     },
+    Hooks {
+        subcommand: Option<HooksSubcommand>,
+    },
     Help,
 }
 
@@ -124,15 +182,43 @@ impl ContextSubcommand {
                                  <black!>--global: Remove specified rules globally</black!>
 
   <em>clear [--global]</em>               <black!>Remove all rules from current profile</black!>
-                                 <black!>--global: Remove global rules</black!>"};
+                                 <black!>--global: Remove global rules</black!>
+
+  <em>hooks</em>                          <black!>View and manage context hooks</black!>"};
     const CLEAR_USAGE: &str = "/context clear [--global]";
+    const HOOKS_AVAILABLE_COMMANDS: &str = color_print::cstr! {"<cyan!>Available subcommands</cyan!>
+  <em>hooks help</em>                         <black!>Show an explanation for context hooks commands</black!>
+
+  <em>hooks add [--global] <<name>></em>        <black!>Add a new command context hook</black!>
+                                         <black!>--global: Add to global hooks</black!>
+         <em>--type <<type>></em>                <black!>Type of hook, valid options: `per_prompt` or `conversation_start`</black!>
+         <em>--command <<command>></em>             <black!>Shell command to execute</black!>
+
+  <em>hooks rm [--global] <<name>></em>         <black!>Remove an existing context hook</black!>
+                                         <black!>--global: Remove from global hooks</black!>
+
+  <em>hooks enable [--global] <<name>></em>     <black!>Enable an existing context hook</black!>
+                                         <black!>--global: Enable in global hooks</black!>
+
+  <em>hooks disable [--global] <<name>></em>    <black!>Disable an existing context hook</black!>
+                                         <black!>--global: Disable in global hooks</black!>
+
+  <em>hooks enable-all [--global]</em>        <black!>Enable all existing context hooks</black!>
+                                         <black!>--global: Enable all in global hooks</black!>
+
+  <em>hooks disable-all [--global]</em>       <black!>Disable all existing context hooks</black!>
+                                         <black!>--global: Disable all in global hooks</black!>"};
     const REMOVE_USAGE: &str = "/context rm [--global] <path1> [path2...]";
     const SHOW_USAGE: &str = "/context show [--expand]";
 
     fn usage_msg(header: impl AsRef<str>) -> String {
         format!("{}\n\n{}", header.as_ref(), Self::AVAILABLE_COMMANDS)
     }
 
+    fn hooks_usage_msg(header: impl AsRef<str>) -> String {
+        format!("{}\n\n{}", header.as_ref(), Self::HOOKS_AVAILABLE_COMMANDS)
+    }
+
     pub fn help_text() -> String {
         color_print::cformat!(
             r#"
@@ -143,6 +229,9 @@ The files matched by these rules provide Amazon Q with additional information
 about your project or environment. Adding relevant files helps Q generate 
 more accurate and helpful responses.
 
+In addition to files, you can specify hooks that will run commands and return 
+the output as context to Amazon Q.
+
 {}
 
 <cyan!>Notes</cyan!>
@@ -154,6 +243,32 @@ more accurate and helpful responses.
             Self::AVAILABLE_COMMANDS
         )
     }
+
+    pub fn hooks_help_text() -> String {
+        color_print::cformat!(
+            r#"
+<magenta,em>(Beta) Context Hooks</magenta,em>
+
+Use context hooks to specify shell commands to run. The output from these 
+commands will be appended to the prompt to Amazon Q. Hooks can be defined 
+in global or local profiles.
+
+<cyan!>Usage: /context hooks [SUBCOMMAND]</cyan!>
+
+<cyan!>Description</cyan!>
+  Show existing global or profile-specific hooks.
+  Alternatively, specify a subcommand to modify the hooks.
+
+{}
+
+<cyan!>Notes</cyan!>
+• Hooks are executed in parallel
+• 'conversation_start' hooks run on the first user prompt and are attached once to the conversation history sent to Amazon Q
+• 'per_prompt' hooks run on each user prompt and are attached to the prompt
+"#,
+            Self::HOOKS_AVAILABLE_COMMANDS
+        )
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -493,6 +608,18 @@ impl Command {
                         "help" => Self::Context {
                             subcommand: ContextSubcommand::Help,
                         },
+                        "hooks" => {
+                            if parts.get(2).is_none() {
+                                return Ok(Self::Context {
+                                    subcommand: ContextSubcommand::Hooks { subcommand: None },
+                                });
+                            };
+
+                            match Self::parse_hooks(&parts) {
+                                Ok(command) => command,
+                                Err(err) => return Err(ContextSubcommand::hooks_usage_msg(err)),
+                            }
+                        },
                         other => {
                             return Err(ContextSubcommand::usage_msg(format!("Unknown subcommand '{}'.", other)));
                         },
@@ -567,6 +694,27 @@ impl Command {
             prompt: input.to_string(),
         })
     }
+
+    // NOTE: Here we use clap to parse the hooks subcommand instead of parsing manually
+    // like the rest of the file.
+    // Since the hooks subcommand has a lot of options, this makes more sense.
+    // Ideally, we parse everything with clap instead of trying to do it manually.
+    fn parse_hooks(parts: &[&str]) -> Result<Self, String> {
+        // Skip the first two parts ("/context" and "hooks")
+        let args = match shlex::split(&parts[1..].join(" ")) {
+            Some(args) => args,
+            None => return Err("Failed to parse arguments".to_string()),
+        };
+
+        // Parse with Clap
+        HooksCommand::try_parse_from(args)
+            .map(|hooks_command| Self::Context {
+                subcommand: ContextSubcommand::Hooks {
+                    subcommand: Some(hooks_command.command),
+                },
+            })
+            .map_err(|e| e.to_string())
+    }
 }
 
 #[cfg(test)]
@@ -688,6 +836,66 @@ mod tests {
             ("/issue \"there was an error in the chat\"", Command::Issue {
                 prompt: Some("\"there was an error in the chat\"".to_string()),
             }),
+            (
+                "/context hooks",
+                context!(ContextSubcommand::Hooks { subcommand: None }),
+            ),
+            (
+                "/context hooks add test --type per_prompt --command 'echo 1' --global",
+                context!(ContextSubcommand::Hooks {
+                    subcommand: Some(HooksSubcommand::Add {
+                        name: "test".to_string(),
+                        global: true,
+                        r#type: "per_prompt".to_string(),
+                        command: "echo 1".to_string()
+                    })
+                }),
+            ),
+            (
+                "/context hooks rm test --global",
+                context!(ContextSubcommand::Hooks {
+                    subcommand: Some(HooksSubcommand::Remove {
+                        name: "test".to_string(),
+                        global: true
+                    })
+                }),
+            ),
+            (
+                "/context hooks enable test --global",
+                context!(ContextSubcommand::Hooks {
+                    subcommand: Some(HooksSubcommand::Enable {
+                        name: "test".to_string(),
+                        global: true
+                    })
+                }),
+            ),
+            (
+                "/context hooks disable test",
+                context!(ContextSubcommand::Hooks {
+                    subcommand: Some(HooksSubcommand::Disable {
+                        name: "test".to_string(),
+                        global: false
+                    })
+                }),
+            ),
+            (
+                "/context hooks enable-all --global",
+                context!(ContextSubcommand::Hooks {
+                    subcommand: Some(HooksSubcommand::EnableAll { global: true })
+                }),
+            ),
+            (
+                "/context hooks disable-all",
+                context!(ContextSubcommand::Hooks {
+                    subcommand: Some(HooksSubcommand::DisableAll { global: false })
+                }),
+            ),
+            (
+                "/context hooks help",
+                context!(ContextSubcommand::Hooks {
+                    subcommand: Some(HooksSubcommand::Help)
+                }),
+            ),
         ];
 
         for (input, parsed) in tests {