feat: recipes can retry with success criteria (#3474)

2025-07-22 10:49:21 +10:00
parent 5f3c7d339c
commit 99cc0a9c81
17 changed files with 1078 additions and 82 deletions
@@ -0,0 +1,498 @@
+use anyhow::Result;
+use std::process::Stdio;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::process::Command;
+use tokio::sync::Mutex;
+use tracing::{debug, info, warn};
+
+use crate::agents::types::SessionConfig;
+use crate::agents::types::{
+    RetryConfig, SuccessCheck, DEFAULT_ON_FAILURE_TIMEOUT_SECONDS, DEFAULT_RETRY_TIMEOUT_SECONDS,
+};
+use crate::config::Config;
+use crate::message::Message;
+use crate::tool_monitor::ToolMonitor;
+
+/// Result of a retry logic evaluation
+#[derive(Debug, Clone, PartialEq)]
+pub enum RetryResult {
+    /// No retry configuration or session available, retry logic skipped
+    Skipped,
+    /// Maximum retry attempts reached, cannot retry further
+    MaxAttemptsReached,
+    /// Success checks passed, no retry needed
+    SuccessChecksPassed,
+    /// Retry is needed and will be performed
+    Retried,
+}
+
+/// Environment variable for configuring retry timeout globally
+const GOOSE_RECIPE_RETRY_TIMEOUT_SECONDS: &str = "GOOSE_RECIPE_RETRY_TIMEOUT_SECONDS";
+
+/// Environment variable for configuring on_failure timeout globally
+const GOOSE_RECIPE_ON_FAILURE_TIMEOUT_SECONDS: &str = "GOOSE_RECIPE_ON_FAILURE_TIMEOUT_SECONDS";
+
+/// Manages retry state and operations for agent execution
+#[derive(Debug)]
+pub struct RetryManager {
+    /// Current number of retry attempts
+    attempts: Arc<Mutex<u32>>,
+    /// Optional tool monitor for reset operations
+    tool_monitor: Option<Arc<Mutex<Option<ToolMonitor>>>>,
+}
+
+impl Default for RetryManager {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl RetryManager {
+    /// Create a new retry manager
+    pub fn new() -> Self {
+        Self {
+            attempts: Arc::new(Mutex::new(0)),
+            tool_monitor: None,
+        }
+    }
+
+    /// Create a new retry manager with tool monitor
+    pub fn with_tool_monitor(tool_monitor: Arc<Mutex<Option<ToolMonitor>>>) -> Self {
+        Self {
+            attempts: Arc::new(Mutex::new(0)),
+            tool_monitor: Some(tool_monitor),
+        }
+    }
+
+    /// Reset the retry attempts counter to 0
+    pub async fn reset_attempts(&self) {
+        let mut attempts = self.attempts.lock().await;
+        *attempts = 0;
+
+        // Reset tool monitor if available
+        if let Some(monitor) = &self.tool_monitor {
+            if let Some(monitor) = monitor.lock().await.as_mut() {
+                monitor.reset();
+            }
+        }
+    }
+
+    /// Increment the retry attempts counter and return the new value
+    pub async fn increment_attempts(&self) -> u32 {
+        let mut attempts = self.attempts.lock().await;
+        *attempts += 1;
+        *attempts
+    }
+
+    /// Get the current retry attempts count
+    pub async fn get_attempts(&self) -> u32 {
+        *self.attempts.lock().await
+    }
+
+    /// Reset status for retry: clear message history and final output tool state
+    async fn reset_status_for_retry(
+        messages: &mut Vec<Message>,
+        initial_messages: &[Message],
+        final_output_tool: &Arc<Mutex<Option<crate::agents::final_output_tool::FinalOutputTool>>>,
+    ) {
+        messages.clear();
+        messages.extend_from_slice(initial_messages);
+        info!("Reset message history to initial state for retry");
+
+        if let Some(final_output_tool) = final_output_tool.lock().await.as_mut() {
+            final_output_tool.final_output = None;
+            info!("Cleared final output tool state for retry");
+        }
+    }
+
+    /// Handle retry logic for the agent reply loop
+    pub async fn handle_retry_logic(
+        &self,
+        messages: &mut Vec<Message>,
+        session: &Option<SessionConfig>,
+        initial_messages: &[Message],
+        final_output_tool: &Arc<Mutex<Option<crate::agents::final_output_tool::FinalOutputTool>>>,
+    ) -> Result<RetryResult> {
+        let Some(session_config) = session else {
+            return Ok(RetryResult::Skipped);
+        };
+
+        let Some(retry_config) = &session_config.retry_config else {
+            return Ok(RetryResult::Skipped);
+        };
+
+        let success = execute_success_checks(&retry_config.checks, retry_config).await?;
+
+        if success {
+            info!("All success checks passed, no retry needed");
+            return Ok(RetryResult::SuccessChecksPassed);
+        }
+
+        let current_attempts = self.get_attempts().await;
+        if current_attempts >= retry_config.max_retries {
+            let error_msg = Message::assistant().with_text(format!(
+                "Maximum retry attempts ({}) exceeded. Unable to complete the task successfully.",
+                retry_config.max_retries
+            ));
+            messages.push(error_msg);
+            warn!(
+                "Maximum retry attempts ({}) exceeded",
+                retry_config.max_retries
+            );
+            return Ok(RetryResult::MaxAttemptsReached);
+        }
+
+        if let Some(on_failure_cmd) = &retry_config.on_failure {
+            info!("Executing on_failure command: {}", on_failure_cmd);
+            execute_on_failure_command(on_failure_cmd, retry_config).await?;
+        }
+
+        Self::reset_status_for_retry(messages, initial_messages, final_output_tool).await;
+
+        let new_attempts = self.increment_attempts().await;
+        info!("Incrementing retry attempts to {}", new_attempts);
+
+        Ok(RetryResult::Retried)
+    }
+}
+
+/// Get the configured timeout duration for retry operations
+/// retry_config.timeout_seconds -> env var -> default
+fn get_retry_timeout(retry_config: &RetryConfig) -> Duration {
+    let timeout_seconds = retry_config
+        .timeout_seconds
+        .or_else(|| {
+            let config = Config::global();
+            config.get_param(GOOSE_RECIPE_RETRY_TIMEOUT_SECONDS).ok()
+        })
+        .unwrap_or(DEFAULT_RETRY_TIMEOUT_SECONDS);
+
+    Duration::from_secs(timeout_seconds)
+}
+
+/// Get the configured timeout duration for on_failure operations
+/// retry_config.on_failure_timeout_seconds -> env var -> default
+fn get_on_failure_timeout(retry_config: &RetryConfig) -> Duration {
+    let timeout_seconds = retry_config
+        .on_failure_timeout_seconds
+        .or_else(|| {
+            let config = Config::global();
+            config
+                .get_param(GOOSE_RECIPE_ON_FAILURE_TIMEOUT_SECONDS)
+                .ok()
+        })
+        .unwrap_or(DEFAULT_ON_FAILURE_TIMEOUT_SECONDS);
+
+    Duration::from_secs(timeout_seconds)
+}
+
+/// Execute all success checks and return true if all pass
+pub async fn execute_success_checks(
+    checks: &[SuccessCheck],
+    retry_config: &RetryConfig,
+) -> Result<bool> {
+    let timeout = get_retry_timeout(retry_config);
+
+    for check in checks {
+        match check {
+            SuccessCheck::Shell { command } => {
+                let result = execute_shell_command(command, timeout).await?;
+                if !result.status.success() {
+                    warn!(
+                        "Success check failed: command '{}' exited with status {}, stderr: {}",
+                        command,
+                        result.status,
+                        String::from_utf8_lossy(&result.stderr)
+                    );
+                    return Ok(false);
+                }
+                info!(
+                    "Success check passed: command '{}' completed successfully",
+                    command
+                );
+            }
+        }
+    }
+    Ok(true)
+}
+
+/// Execute a shell command with cross-platform compatibility and mandatory timeout
+pub async fn execute_shell_command(
+    command: &str,
+    timeout: std::time::Duration,
+) -> Result<std::process::Output> {
+    debug!(
+        "Executing shell command with timeout {:?}: {}",
+        timeout, command
+    );
+
+    let future = async {
+        let mut cmd = if cfg!(target_os = "windows") {
+            let mut cmd = Command::new("cmd");
+            cmd.args(["/C", command]);
+            cmd
+        } else {
+            let mut cmd = Command::new("sh");
+            cmd.args(["-c", command]);
+            cmd
+        };
+
+        let output = cmd
+            .stdout(Stdio::piped())
+            .stderr(Stdio::piped())
+            .stdin(Stdio::null())
+            .kill_on_drop(true)
+            .output()
+            .await?;
+
+        debug!(
+            "Shell command completed with status: {}, stdout: {}, stderr: {}",
+            output.status,
+            String::from_utf8_lossy(&output.stdout),
+            String::from_utf8_lossy(&output.stderr)
+        );
+
+        Ok(output)
+    };
+
+    match tokio::time::timeout(timeout, future).await {
+        Ok(result) => result,
+        Err(_) => {
+            let error_msg = format!("Shell command timed out after {:?}: {}", timeout, command);
+            warn!("{}", error_msg);
+            Err(anyhow::anyhow!("{}", error_msg))
+        }
+    }
+}
+
+/// Execute an on_failure command and return an error if it fails
+pub async fn execute_on_failure_command(command: &str, retry_config: &RetryConfig) -> Result<()> {
+    let timeout = get_on_failure_timeout(retry_config);
+    info!(
+        "Executing on_failure command with timeout {:?}: {}",
+        timeout, command
+    );
+
+    let output = match execute_shell_command(command, timeout).await {
+        Ok(output) => output,
+        Err(e) => {
+            if e.to_string().contains("timed out") {
+                let error_msg = format!(
+                    "On_failure command timed out after {:?}: {}",
+                    timeout, command
+                );
+                warn!("{}", error_msg);
+                return Err(anyhow::anyhow!(error_msg));
+            } else {
+                warn!("On_failure command execution error: {}", e);
+                return Err(e);
+            }
+        }
+    };
+
+    if !output.status.success() {
+        let error_msg = format!(
+            "On_failure command failed: command '{}' exited with status {}, stderr: {}",
+            command,
+            output.status,
+            String::from_utf8_lossy(&output.stderr)
+        );
+        warn!("{}", error_msg);
+        return Err(anyhow::anyhow!(error_msg));
+    } else {
+        info!("On_failure command completed successfully: {}", command);
+    }
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::agents::types::SuccessCheck;
+
+    fn create_test_retry_config() -> RetryConfig {
+        RetryConfig {
+            max_retries: 3,
+            checks: vec![],
+            on_failure: None,
+            timeout_seconds: Some(60),
+            on_failure_timeout_seconds: Some(120),
+        }
+    }
+
+    #[test]
+    fn test_retry_result_enum() {
+        assert_ne!(RetryResult::Skipped, RetryResult::MaxAttemptsReached);
+        assert_ne!(RetryResult::Skipped, RetryResult::SuccessChecksPassed);
+        assert_ne!(RetryResult::Skipped, RetryResult::Retried);
+        assert_ne!(
+            RetryResult::MaxAttemptsReached,
+            RetryResult::SuccessChecksPassed
+        );
+        assert_ne!(RetryResult::MaxAttemptsReached, RetryResult::Retried);
+        assert_ne!(RetryResult::SuccessChecksPassed, RetryResult::Retried);
+
+        let result = RetryResult::Retried;
+        let cloned = result.clone();
+        assert_eq!(result, cloned);
+
+        let debug_str = format!("{:?}", RetryResult::MaxAttemptsReached);
+        assert!(debug_str.contains("MaxAttemptsReached"));
+    }
+
+    #[tokio::test]
+    async fn test_execute_success_checks_all_pass() {
+        let checks = vec![
+            SuccessCheck::Shell {
+                command: "echo 'test'".to_string(),
+            },
+            SuccessCheck::Shell {
+                command: "true".to_string(),
+            },
+        ];
+        let retry_config = create_test_retry_config();
+
+        let result = execute_success_checks(&checks, &retry_config).await;
+        assert!(result.is_ok());
+        assert!(result.unwrap());
+    }
+
+    #[tokio::test]
+    async fn test_execute_success_checks_one_fails() {
+        let checks = vec![
+            SuccessCheck::Shell {
+                command: "echo 'test'".to_string(),
+            },
+            SuccessCheck::Shell {
+                command: "false".to_string(),
+            },
+        ];
+        let retry_config = create_test_retry_config();
+
+        let result = execute_success_checks(&checks, &retry_config).await;
+        assert!(result.is_ok());
+        assert!(!result.unwrap());
+    }
+
+    #[tokio::test]
+    async fn test_execute_shell_command_success() {
+        let result = execute_shell_command("echo 'hello world'", Duration::from_secs(30)).await;
+        assert!(result.is_ok());
+        let output = result.unwrap();
+        assert!(output.status.success());
+        assert!(String::from_utf8_lossy(&output.stdout).contains("hello world"));
+    }
+
+    #[tokio::test]
+    async fn test_execute_shell_command_failure() {
+        let result = execute_shell_command("false", Duration::from_secs(30)).await;
+        assert!(result.is_ok());
+        let output = result.unwrap();
+        assert!(!output.status.success());
+    }
+
+    #[tokio::test]
+    async fn test_execute_on_failure_command_success() {
+        let retry_config = create_test_retry_config();
+        let result = execute_on_failure_command("echo 'cleanup'", &retry_config).await;
+        assert!(result.is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_execute_on_failure_command_failure() {
+        let retry_config = create_test_retry_config();
+        let result = execute_on_failure_command("false", &retry_config).await;
+        assert!(result.is_err());
+    }
+
+    #[tokio::test]
+    async fn test_shell_command_timeout() {
+        let timeout = std::time::Duration::from_millis(100);
+        let result = if cfg!(target_os = "windows") {
+            execute_shell_command("timeout /t 1", timeout).await
+        } else {
+            execute_shell_command("sleep 1", timeout).await
+        };
+
+        assert!(result.is_err());
+    }
+
+    #[tokio::test]
+    async fn test_get_retry_timeout_uses_config_default() {
+        let retry_config = RetryConfig {
+            max_retries: 1,
+            checks: vec![],
+            on_failure: None,
+            timeout_seconds: None,
+            on_failure_timeout_seconds: None,
+        };
+
+        let timeout = get_retry_timeout(&retry_config);
+        assert_eq!(timeout, Duration::from_secs(DEFAULT_RETRY_TIMEOUT_SECONDS));
+    }
+
+    #[tokio::test]
+    async fn test_get_retry_timeout_uses_retry_config() {
+        let retry_config = RetryConfig {
+            max_retries: 1,
+            checks: vec![],
+            on_failure: None,
+            timeout_seconds: Some(120),
+            on_failure_timeout_seconds: None,
+        };
+
+        let timeout = get_retry_timeout(&retry_config);
+        assert_eq!(timeout, Duration::from_secs(120));
+    }
+
+    #[tokio::test]
+    async fn test_get_on_failure_timeout_uses_config_default() {
+        let retry_config = RetryConfig {
+            max_retries: 1,
+            checks: vec![],
+            on_failure: None,
+            timeout_seconds: None,
+            on_failure_timeout_seconds: None,
+        };
+
+        let timeout = get_on_failure_timeout(&retry_config);
+        assert_eq!(
+            timeout,
+            Duration::from_secs(DEFAULT_ON_FAILURE_TIMEOUT_SECONDS)
+        );
+    }
+
+    #[tokio::test]
+    async fn test_get_on_failure_timeout_uses_retry_config() {
+        let retry_config = RetryConfig {
+            max_retries: 1,
+            checks: vec![],
+            on_failure: None,
+            timeout_seconds: None,
+            on_failure_timeout_seconds: Some(900),
+        };
+
+        let timeout = get_on_failure_timeout(&retry_config);
+        assert_eq!(timeout, Duration::from_secs(900));
+    }
+
+    #[tokio::test]
+    async fn test_on_failure_timeout_different_from_retry_timeout() {
+        let retry_config = RetryConfig {
+            max_retries: 1,
+            checks: vec![],
+            on_failure: None,
+            timeout_seconds: Some(60),
+            on_failure_timeout_seconds: Some(300),
+        };
+
+        let retry_timeout = get_retry_timeout(&retry_config);
+        let on_failure_timeout = get_on_failure_timeout(&retry_config);
+
+        assert_eq!(retry_timeout, Duration::from_secs(60));
+        assert_eq!(on_failure_timeout, Duration::from_secs(300));
+        assert_ne!(retry_timeout, on_failure_timeout);
+    }
+}