Merge pull request #522 from stakpak/feat/trigger-check-customization

Add customizable check logic & prompt improvements

Author: kajogo777

cli/src/commands/watch/commands/init.rs

@@ -59,6 +59,12 @@ timeout = "30m"
 # Maximum time for check scripts (default: 30s)
 check_timeout = "30s"
 
+# Determines which check script exit codes trigger the agent (default: "success")
+# - "success": trigger on exit 0
+# - "failure": trigger on non-zero exit codes (1+)
+# - "any": trigger regardless of exit code (only timeout/error prevents trigger)
+# check_trigger_on = "success"
+
 # Enable Slack tools for agent (experimental, default: false)
 # enable_slack_tools = false
 
@@ -85,9 +91,14 @@ Output a summary report but DO NOT make any changes to the system.
 This is a read-only health check.
 """
 # Optional: check script that determines if agent should run
-# Exit 0 = run agent, Exit 1 = skip, Exit 2+ = error
 # check = "~/.stakpak/triggers/check-weekday.sh"
 
+# Optional: which exit codes trigger the agent (overrides default)
+# - "success": trigger on exit 0 (default)
+# - "failure": trigger on non-zero exit codes (1+)
+# - "any": trigger regardless of exit code
+# check_trigger_on = "success"
+
 # Optional: override default timeout for this trigger
 # timeout = "10m"
 
@@ -112,6 +123,20 @@ This is a read-only health check.
 # DO NOT make any changes, just report findings.
 # """
 
+# Example: Alert on service failure (trigger on failure)
+# [[triggers]]
+# name = "service-down-alert"
+# schedule = "*/5 * * * *"  # Every 5 minutes
+# check = "~/.stakpak/triggers/check-service-health.sh"
+# check_trigger_on = "failure"  # Only wake agent when check fails
+# prompt = """
+# The service health check failed. Investigate and report:
+# - Which services are down
+# - Recent error logs
+# - Possible root causes
+# DO NOT restart services automatically, just report findings.
+# """
+
 # Example: Security advisory check (read-only)
 # [[triggers]]
 # name = "security-advisories"

cli/src/commands/watch/commands/run.rs

@@ -278,38 +278,39 @@ async fn handle_trigger_event(
                     return Ok(());
                 }
 
-                if result.failed() {
-                    warn!(
+                // Determine if we should trigger based on check_trigger_on setting
+                let exit_code = result.exit_code.unwrap_or(-1);
+                let check_trigger_on = trigger.effective_check_trigger_on(&config.defaults);
+                let should_trigger = check_trigger_on.should_trigger(exit_code);
+
+                if !should_trigger {
+                    info!(
                         trigger = %trigger.name,
-                        exit_code = ?result.exit_code,
-                        "Check script failed"
+                        exit_code = exit_code,
+                        check_trigger_on = %check_trigger_on,
+                        "Check script did not meet trigger condition"
                     );
                     print_event(
-                        "fail",
+                        "skip",
                         &trigger.name,
-                        &format!("Check failed (exit {})", result.exit_code.unwrap_or(-1)),
+                        &format!(
+                            "Skipped (exit {} does not match trigger_on={})",
+                            exit_code, check_trigger_on
+                        ),
                     );
-                    db.update_run_finished(
-                        run_id,
-                        RunStatus::Failed,
-                        Some("Check script failed"),
-                        None,
-                        None,
-                    )
-                    .await
-                    .map_err(|e| format!("Failed to update run status: {}", e))?;
-                    return Ok(());
-                }
-
-                if result.skipped() {
-                    info!(trigger = %trigger.name, "Check script returned skip (exit 1)");
-                    print_event("skip", &trigger.name, "Skipped (check returned exit 1)");
                     db.update_run_finished(run_id, RunStatus::Skipped, None, None, None)
                         .await
                         .map_err(|e| format!("Failed to update run status: {}", e))?;
                     return Ok(());
                 }
 
+                info!(
+                    trigger = %trigger.name,
+                    exit_code = exit_code,
+                    check_trigger_on = %check_trigger_on,
+                    "Check script met trigger condition"
+                );
+
                 Some(result)
             }
             Err(e) => {

cli/src/commands/watch/commands/trigger.rs

@@ -147,12 +147,21 @@ pub async fn fire_trigger(name: &str, dry_run: bool) -> Result<(), String> {
 
                     if result.timed_out {
                         println!("\n\x1b[31mCheck script timed out\x1b[0m");
-                    } else if result.failed() {
-                        println!("\n\x1b[31mCheck script failed (exit {})\x1b[0m", exit_code);
-                    } else if result.skipped() {
-                        println!(
-                            "\n\x1b[33mCheck script returned skip (exit 1) - agent would not be woken\x1b[0m"
-                        );
+                    } else {
+                        let check_trigger_on = trigger.effective_check_trigger_on(&config.defaults);
+                        let should_trigger = check_trigger_on.should_trigger(exit_code);
+                        println!("Check trigger_on: {}", check_trigger_on);
+                        if should_trigger {
+                            println!(
+                                "\n\x1b[32mCheck passed (exit {} matches trigger_on={}) - agent would be woken\x1b[0m",
+                                exit_code, check_trigger_on
+                            );
+                        } else {
+                            println!(
+                                "\n\x1b[33mCheck skipped (exit {} does not match trigger_on={}) - agent would not be woken\x1b[0m",
+                                exit_code, check_trigger_on
+                            );
+                        }
                     }
 
                     Some(result)

cli/src/commands/watch/config.rs

@@ -57,6 +57,40 @@ fn default_log_dir() -> String {
     "~/.stakpak/watch/logs".to_string()
 }
 
+/// Determines which check script exit codes trigger the agent.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum CheckTriggerOn {
+    /// Trigger agent only on exit code 0 (default behavior).
+    #[default]
+    Success,
+    /// Trigger agent on any non-zero exit code (1+).
+    Failure,
+    /// Trigger agent regardless of exit code (only timeout/error prevents trigger).
+    Any,
+}
+
+impl CheckTriggerOn {
+    /// Returns true if the given exit code should trigger the agent.
+    pub fn should_trigger(&self, exit_code: i32) -> bool {
+        match self {
+            CheckTriggerOn::Success => exit_code == 0,
+            CheckTriggerOn::Failure => exit_code != 0,
+            CheckTriggerOn::Any => true,
+        }
+    }
+}
+
+impl std::fmt::Display for CheckTriggerOn {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            CheckTriggerOn::Success => write!(f, "success"),
+            CheckTriggerOn::Failure => write!(f, "failure"),
+            CheckTriggerOn::Any => write!(f, "any"),
+        }
+    }
+}
+
 /// Default values applied to triggers when not specified.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct WatchDefaults {
@@ -84,6 +118,13 @@ pub struct WatchDefaults {
     /// When true, the agent will pause and exit with code 10 when tools need approval.
     #[serde(default = "default_pause_on_approval")]
     pub pause_on_approval: bool,
+
+    /// Determines which check script exit codes trigger the agent.
+    /// - "success" (default): trigger on exit 0
+    /// - "failure": trigger on non-zero exit codes (1+)
+    /// - "any": trigger regardless of exit code
+    #[serde(default)]
+    pub check_trigger_on: CheckTriggerOn,
 }
 
 impl Default for WatchDefaults {
@@ -95,6 +136,7 @@ impl Default for WatchDefaults {
             enable_slack_tools: false,
             enable_subagents: false,
             pause_on_approval: default_pause_on_approval(),
+            check_trigger_on: CheckTriggerOn::default(),
         }
     }
 }
@@ -133,6 +175,13 @@ pub struct Trigger {
     #[serde(default, with = "option_humantime_serde")]
     pub check_timeout: Option<Duration>,
 
+    /// Determines which check script exit codes trigger the agent.
+    /// Falls back to defaults.check_trigger_on if not specified.
+    /// - "success" (default): trigger on exit 0
+    /// - "failure": trigger on non-zero exit codes (1+)
+    /// - "any": trigger regardless of exit code
+    pub check_trigger_on: Option<CheckTriggerOn>,
+
     /// Prompt to pass to the agent when triggered.
     pub prompt: String,
 
@@ -177,6 +226,11 @@ impl Trigger {
         self.check_timeout.unwrap_or(defaults.check_timeout)
     }
 
+    /// Get the effective check_trigger_on, falling back to defaults.
+    pub fn effective_check_trigger_on(&self, defaults: &WatchDefaults) -> CheckTriggerOn {
+        self.check_trigger_on.unwrap_or(defaults.check_trigger_on)
+    }
+
     /// Get the effective enable_slack_tools, falling back to defaults.
     pub fn effective_enable_slack_tools(&self, defaults: &WatchDefaults) -> bool {
         self.enable_slack_tools
@@ -627,6 +681,134 @@ prompt = "Test"
 [[triggers]]
 name = "test"
 schedule = "0 * * * *"
+"#;
+
+        let result = WatchConfig::parse(config_str);
+        assert!(result.is_err());
+        assert!(matches!(result.unwrap_err(), ConfigError::ParseError(_)));
+    }
+
+    #[test]
+    fn test_check_trigger_on_should_trigger() {
+        // Success mode: only exit 0 triggers
+        assert!(CheckTriggerOn::Success.should_trigger(0));
+        assert!(!CheckTriggerOn::Success.should_trigger(1));
+        assert!(!CheckTriggerOn::Success.should_trigger(2));
+        assert!(!CheckTriggerOn::Success.should_trigger(-1));
+
+        // Failure mode: any non-zero triggers
+        assert!(!CheckTriggerOn::Failure.should_trigger(0));
+        assert!(CheckTriggerOn::Failure.should_trigger(1));
+        assert!(CheckTriggerOn::Failure.should_trigger(2));
+        assert!(CheckTriggerOn::Failure.should_trigger(-1));
+
+        // Any mode: all exit codes trigger
+        assert!(CheckTriggerOn::Any.should_trigger(0));
+        assert!(CheckTriggerOn::Any.should_trigger(1));
+        assert!(CheckTriggerOn::Any.should_trigger(2));
+        assert!(CheckTriggerOn::Any.should_trigger(-1));
+    }
+
+    #[test]
+    fn test_check_trigger_on_default() {
+        assert_eq!(CheckTriggerOn::default(), CheckTriggerOn::Success);
+    }
+
+    #[test]
+    fn test_check_trigger_on_display() {
+        assert_eq!(CheckTriggerOn::Success.to_string(), "success");
+        assert_eq!(CheckTriggerOn::Failure.to_string(), "failure");
+        assert_eq!(CheckTriggerOn::Any.to_string(), "any");
+    }
+
+    #[test]
+    fn test_check_trigger_on_parsing() {
+        // Test parsing from TOML
+        let config_str = r#"
+[[triggers]]
+name = "success-trigger"
+schedule = "0 * * * *"
+prompt = "Test"
+check_trigger_on = "success"
+
+[[triggers]]
+name = "failure-trigger"
+schedule = "0 * * * *"
+prompt = "Test"
+check_trigger_on = "failure"
+
+[[triggers]]
+name = "any-trigger"
+schedule = "0 * * * *"
+prompt = "Test"
+check_trigger_on = "any"
+
+[[triggers]]
+name = "default-trigger"
+schedule = "0 * * * *"
+prompt = "Test"
+"#;
+
+        let config = WatchConfig::parse(config_str).expect("Should parse check_trigger_on values");
+        assert_eq!(config.triggers.len(), 4);
+
+        assert_eq!(
+            config.triggers[0].check_trigger_on,
+            Some(CheckTriggerOn::Success)
+        );
+        assert_eq!(
+            config.triggers[1].check_trigger_on,
+            Some(CheckTriggerOn::Failure)
+        );
+        assert_eq!(
+            config.triggers[2].check_trigger_on,
+            Some(CheckTriggerOn::Any)
+        );
+        assert_eq!(config.triggers[3].check_trigger_on, None);
+    }
+
+    #[test]
+    fn test_check_trigger_on_defaults_fallback() {
+        let config_str = r#"
+[defaults]
+check_trigger_on = "failure"
+
+[[triggers]]
+name = "uses-default"
+schedule = "0 * * * *"
+prompt = "Test"
+
+[[triggers]]
+name = "overrides-default"
+schedule = "0 * * * *"
+prompt = "Test"
+check_trigger_on = "success"
+"#;
+
+        let config =
+            WatchConfig::parse(config_str).expect("Should parse check_trigger_on with defaults");
+
+        // First trigger should use default (failure)
+        assert_eq!(
+            config.triggers[0].effective_check_trigger_on(&config.defaults),
+            CheckTriggerOn::Failure
+        );
+
+        // Second trigger should override to success
+        assert_eq!(
+            config.triggers[1].effective_check_trigger_on(&config.defaults),
+            CheckTriggerOn::Success
+        );
+    }
+
+    #[test]
+    fn test_check_trigger_on_invalid_value() {
+        let config_str = r#"
+[[triggers]]
+name = "invalid"
+schedule = "0 * * * *"
+prompt = "Test"
+check_trigger_on = "invalid"
 "#;
 
         let result = WatchConfig::parse(config_str);