feat: added env variable USER_PROMPT and enabled removal + sanitizing of en… (#2244)

Author: syaamkhan

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

@@ -164,6 +164,7 @@ impl HookExecutor {
         &mut self,
         hooks: Vec<&Hook>,
         output: &mut impl Write,
+        prompt: Option<&str>
     ) -> Result<Vec<(Hook, String)>, ChatError> {
         let mut results = Vec::with_capacity(hooks.len());
         let mut futures = FuturesUnordered::new();
@@ -181,7 +182,7 @@ impl HookExecutor {
                 results.push((index, (hook.clone(), cached.clone())));
                 continue;
             }
-            let future = self.execute_hook(hook);
+            let future = self.execute_hook(hook, prompt);
             futures.push(async move { (index, future.await) });
         }
 
@@ -299,38 +300,47 @@ impl HookExecutor {
         Ok(results.into_iter().map(|(_, r)| r).collect())
     }
 
-    async fn execute_hook<'a>(&self, hook: &'a Hook) -> (&'a Hook, Result<String>, Duration) {
+    async fn execute_hook<'a>(&self, hook: &'a Hook, prompt: Option<&str>) -> (&'a Hook, Result<String>, Duration) {
         let start_time = Instant::now();
         let result = match hook.r#type {
-            HookType::Inline => self.execute_inline_hook(hook).await,
+            HookType::Inline => self.execute_inline_hook(hook, prompt).await,
         };
 
         (hook, result, start_time.elapsed())
     }
 
-    async fn execute_inline_hook(&self, hook: &Hook) -> Result<String> {
+    async fn execute_inline_hook(&self, hook: &Hook, user_prompt: Option<&str>) -> Result<String> {
         let command = hook.command.as_ref().ok_or_else(|| eyre!("no command specified"))?;
 
         #[cfg(unix)]
-        let command_future = tokio::process::Command::new("bash")
-            .arg("-c")
+        let mut cmd = tokio::process::Command::new("bash");
+        #[cfg(unix)]
+        let cmd = cmd.arg("-c")
             .arg(command)
             .stdin(Stdio::piped())
             .stdout(Stdio::piped())
-            .stderr(Stdio::piped())
-            .output();
+            .stderr(Stdio::piped());
 
         #[cfg(windows)]
-        let command_future = tokio::process::Command::new("cmd")
-            .arg("/C")
+        let mut cmd = tokio::process::Command::new("cmd");
+        #[cfg(windows)]
+        let cmd = cmd.arg("/C")
             .arg(command)
             .stdin(Stdio::piped())
             .stdout(Stdio::piped())
-            .stderr(Stdio::piped())
-            .output();
+            .stderr(Stdio::piped());
 
         let timeout = Duration::from_millis(hook.timeout_ms);
 
+        // Set USER_PROMPT environment variable if provided
+        if let Some(prompt) = user_prompt {
+            // Sanitize the prompt to avoid issues with special characters
+            let sanitized_prompt = sanitize_user_prompt(prompt);
+            cmd.env("USER_PROMPT", sanitized_prompt);
+        }
+
+        let command_future = cmd.output();
+
         // Run with timeout
         match tokio::time::timeout(timeout, command_future).await {
             Ok(result) => {
@@ -387,6 +397,19 @@ impl HookExecutor {
     }
 }
 
+/// Sanitizes a string value to be used as an environment variable
+fn sanitize_user_prompt(input: &str) -> String {
+    // Limit the size of input to first 4096 characters
+    let truncated = if input.len() > 4096 {
+        &input[0..4096]
+    } else {
+        input
+    };
+    
+    // Remove any potentially problematic characters
+    truncated.replace(|c: char| c.is_control() && c != '\n' && c != '\r' && c != '\t', "")
+}
+
 #[deny(missing_docs)]
 #[derive(Debug, PartialEq, Args)]
 #[command(

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

@@ -272,7 +272,7 @@ impl ContextManager {
     /// Run all the currently enabled hooks from both the global and profile contexts.
     /// # Returns
     /// A vector containing pairs of a [`Hook`] definition and its execution output
-    pub async fn run_hooks(&mut self, output: &mut impl Write) -> Result<Vec<(Hook, String)>, ChatError> {
+    pub async fn run_hooks(&mut self, output: &mut impl Write, prompt: Option<&str>) -> Result<Vec<(Hook, String)>, ChatError> {
         let hooks = self
             .profile_config
             .hooks
@@ -282,7 +282,7 @@ impl ContextManager {
                 hook as &Hook
             })
             .collect::<Vec<_>>();
-        self.hook_executor.run_hooks(hooks, output).await
+        self.hook_executor.run_hooks(hooks, output, prompt).await
     }
 }
 

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

@@ -360,7 +360,10 @@ impl ConversationState {
         // Run hooks and add to conversation start and next user message.
         let mut conversation_start_context = None;
         if let (true, Some(cm)) = (run_hooks, self.context_manager.as_mut()) {
-            let hook_results = cm.run_hooks(output).await?;
+            // Get the user prompt from next_message if available
+            let user_prompt = self.next_message.as_ref().and_then(|m| m.prompt());
+            let hook_results = cm.run_hooks(output, user_prompt).await?;
+
             conversation_start_context = Some(format_hook_context(hook_results.iter(), HookTrigger::ConversationStart));
 
             // add per prompt content to next_user_message if available