owlen/crates/platform/hooks/src/lib.rs

use color_eyre::eyre::{Result, eyre};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::path::PathBuf;
use std::process::Stdio;
use tokio::io::AsyncWriteExt;
use tokio::process::Command;
use tokio::time::timeout;
use std::time::Duration;

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "event", rename_all = "camelCase")]
pub enum HookEvent {
    #[serde(rename_all = "camelCase")]
    PreToolUse {
        tool: String,
        args: Value,
    },
    #[serde(rename_all = "camelCase")]
    PostToolUse {
        tool: String,
        result: Value,
    },
    #[serde(rename_all = "camelCase")]
    SessionStart {
        session_id: String,
    },
    #[serde(rename_all = "camelCase")]
    SessionEnd {
        session_id: String,
    },
    #[serde(rename_all = "camelCase")]
    UserPromptSubmit {
        prompt: String,
    },
    PreCompact,
    /// Called before the agent stops - allows validation of completion
    #[serde(rename_all = "camelCase")]
    Stop {
        /// Reason for stopping (e.g., "task_complete", "max_iterations", "user_interrupt")
        reason: String,
        /// Number of messages in conversation
        num_messages: usize,
        /// Number of tool calls made
        num_tool_calls: usize,
    },
    /// Called before a subagent stops
    #[serde(rename_all = "camelCase")]
    SubagentStop {
        /// Unique identifier for the subagent
        agent_id: String,
        /// Type of subagent (e.g., "explore", "code-reviewer")
        agent_type: String,
        /// Reason for stopping
        reason: String,
    },
    /// Called when a notification is sent to the user
    #[serde(rename_all = "camelCase")]
    Notification {
        /// Notification message
        message: String,
        /// Notification type (e.g., "info", "warning", "error")
        notification_type: String,
    },
}

impl HookEvent {
    /// Get the hook name for this event (used to find the hook script)
    pub fn hook_name(&self) -> &str {
        match self {
            HookEvent::PreToolUse { .. } => "PreToolUse",
            HookEvent::PostToolUse { .. } => "PostToolUse",
            HookEvent::SessionStart { .. } => "SessionStart",
            HookEvent::SessionEnd { .. } => "SessionEnd",
            HookEvent::UserPromptSubmit { .. } => "UserPromptSubmit",
            HookEvent::PreCompact => "PreCompact",
            HookEvent::Stop { .. } => "Stop",
            HookEvent::SubagentStop { .. } => "SubagentStop",
            HookEvent::Notification { .. } => "Notification",
        }
    }
}

/// Simple hook result for backwards compatibility
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum HookResult {
    Allow,
    Deny,
}

/// Extended hook output with additional control options
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct HookOutput {
    /// Whether to continue execution (default: true if exit code 0)
    #[serde(default = "default_continue")]
    pub continue_execution: bool,
    /// Whether to suppress showing the result to the user
    #[serde(default)]
    pub suppress_output: bool,
    /// System message to inject into the conversation
    #[serde(default)]
    pub system_message: Option<String>,
    /// Permission decision override
    #[serde(default)]
    pub permission_decision: Option<HookPermission>,
    /// Modified input/args for the tool (PreToolUse only)
    #[serde(default)]
    pub updated_input: Option<Value>,
}

impl Default for HookOutput {
    fn default() -> Self {
        Self {
            continue_execution: true,
            suppress_output: false,
            system_message: None,
            permission_decision: None,
            updated_input: None,
        }
    }
}

fn default_continue() -> bool {
    true
}

/// Permission decision from a hook
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum HookPermission {
    Allow,
    Deny,
    Ask,
}

impl HookOutput {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn allow() -> Self {
        Self {
            continue_execution: true,
            ..Default::default()
        }
    }

    pub fn deny() -> Self {
        Self {
            continue_execution: false,
            ..Default::default()
        }
    }

    pub fn with_system_message(mut self, message: impl Into<String>) -> Self {
        self.system_message = Some(message.into());
        self
    }

    pub fn with_permission(mut self, permission: HookPermission) -> Self {
        self.permission_decision = Some(permission);
        self
    }

    /// Convert to simple HookResult for backwards compatibility
    pub fn to_result(&self) -> HookResult {
        if self.continue_execution {
            HookResult::Allow
        } else {
            HookResult::Deny
        }
    }
}

/// A registered hook that can be executed
#[derive(Debug, Clone)]
struct Hook {
    event: String,  // Event name like "PreToolUse", "PostToolUse", etc.
    command: String,
    pattern: Option<String>,  // Optional regex pattern for matching tool names
    timeout: Option<u64>,
}

pub struct HookManager {
    project_root: PathBuf,
    hooks: Vec<Hook>,
}

impl HookManager {
    pub fn new(project_root: &str) -> Self {
        Self {
            project_root: PathBuf::from(project_root),
            hooks: Vec::new(),
        }
    }

    /// Register a single hook
    pub fn register_hook(&mut self, event: String, command: String, pattern: Option<String>, timeout: Option<u64>) {
        self.hooks.push(Hook {
            event,
            command,
            pattern,
            timeout,
        });
    }

    /// Execute a hook for the given event
    ///
    /// Returns:
    /// - Ok(HookResult::Allow) if hook succeeds or doesn't exist (exit code 0 or no hook)
    /// - Ok(HookResult::Deny) if hook denies (exit code 2)
    /// - Err if hook fails (other exit codes) or times out
    pub async fn execute(&self, event: &HookEvent, timeout_ms: Option<u64>) -> Result<HookResult> {
        // First check for legacy file-based hooks
        let hook_path = self.get_hook_path(event);
        let has_file_hook = hook_path.exists();

        // Get registered hooks for this event
        let event_name = event.hook_name();
        let mut matching_hooks: Vec<&Hook> = self.hooks.iter()
            .filter(|h| h.event == event_name)
            .collect();

        // If we need to filter by pattern (for PreToolUse events)
        if let HookEvent::PreToolUse { tool, .. } = event {
            matching_hooks.retain(|h| {
                if let Some(pattern) = &h.pattern {
                    // Use regex to match tool name
                    if let Ok(re) = regex::Regex::new(pattern) {
                        re.is_match(tool)
                    } else {
                        false
                    }
                } else {
                    true  // No pattern means match all
                }
            });
        }

        // If no hooks at all, allow by default
        if !has_file_hook && matching_hooks.is_empty() {
            return Ok(HookResult::Allow);
        }

        // Execute file-based hook first (if exists)
        if has_file_hook {
            let result = self.execute_hook_command(&hook_path.to_string_lossy(), event, timeout_ms).await?;
            if result == HookResult::Deny {
                return Ok(HookResult::Deny);
            }
        }

        // Execute registered hooks
        for hook in matching_hooks {
            let hook_timeout = hook.timeout.or(timeout_ms);
            let result = self.execute_hook_command(&hook.command, event, hook_timeout).await?;
            if result == HookResult::Deny {
                return Ok(HookResult::Deny);
            }
        }

        Ok(HookResult::Allow)
    }

    /// Execute a single hook command
    async fn execute_hook_command(&self, command: &str, event: &HookEvent, timeout_ms: Option<u64>) -> Result<HookResult> {
        // Serialize event to JSON
        let input_json = serde_json::to_string(event)?;

        // Spawn the hook process
        let mut child = Command::new("sh")
            .arg("-c")
            .arg(command)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .current_dir(&self.project_root)
            .spawn()?;

        // Write JSON input to stdin
        if let Some(mut stdin) = child.stdin.take() {
            stdin.write_all(input_json.as_bytes()).await?;
            stdin.flush().await?;
            drop(stdin); // Close stdin
        }

        // Wait for process with timeout
        let result = if let Some(ms) = timeout_ms {
            timeout(Duration::from_millis(ms), child.wait_with_output()).await
        } else {
            Ok(child.wait_with_output().await)
        };

        match result {
            Ok(Ok(output)) => {
                // Check exit code
                match output.status.code() {
                    Some(0) => Ok(HookResult::Allow),
                    Some(2) => Ok(HookResult::Deny),
                    Some(code) => Err(eyre!(
                        "Hook {} failed with exit code {}: {}",
                        event.hook_name(),
                        code,
                        String::from_utf8_lossy(&output.stderr)
                    )),
                    None => Err(eyre!("Hook {} terminated by signal", event.hook_name())),
                }
            }
            Ok(Err(e)) => Err(eyre!("Failed to execute hook {}: {}", event.hook_name(), e)),
            Err(_) => Err(eyre!("Hook {} timed out", event.hook_name())),
        }
    }

    /// Execute a hook and return extended output
    ///
    /// This method parses JSON output from stdout if the hook provides it,
    /// otherwise falls back to exit code interpretation.
    pub async fn execute_extended(&self, event: &HookEvent, timeout_ms: Option<u64>) -> Result<HookOutput> {
        // First check for legacy file-based hooks
        let hook_path = self.get_hook_path(event);
        let has_file_hook = hook_path.exists();

        // Get registered hooks for this event
        let event_name = event.hook_name();
        let mut matching_hooks: Vec<&Hook> = self.hooks.iter()
            .filter(|h| h.event == event_name)
            .collect();

        // If we need to filter by pattern (for PreToolUse events)
        if let HookEvent::PreToolUse { tool, .. } = event {
            matching_hooks.retain(|h| {
                if let Some(pattern) = &h.pattern {
                    if let Ok(re) = regex::Regex::new(pattern) {
                        re.is_match(tool)
                    } else {
                        false
                    }
                } else {
                    true
                }
            });
        }

        // If no hooks at all, allow by default
        if !has_file_hook && matching_hooks.is_empty() {
            return Ok(HookOutput::allow());
        }

        let mut combined_output = HookOutput::allow();

        // Execute file-based hook first (if exists)
        if has_file_hook {
            let output = self.execute_hook_extended(&hook_path.to_string_lossy(), event, timeout_ms).await?;
            combined_output = Self::merge_outputs(combined_output, output);
            if !combined_output.continue_execution {
                return Ok(combined_output);
            }
        }

        // Execute registered hooks
        for hook in matching_hooks {
            let hook_timeout = hook.timeout.or(timeout_ms);
            let output = self.execute_hook_extended(&hook.command, event, hook_timeout).await?;
            combined_output = Self::merge_outputs(combined_output, output);
            if !combined_output.continue_execution {
                return Ok(combined_output);
            }
        }

        Ok(combined_output)
    }

    /// Execute a single hook command and return extended output
    async fn execute_hook_extended(&self, command: &str, event: &HookEvent, timeout_ms: Option<u64>) -> Result<HookOutput> {
        let input_json = serde_json::to_string(event)?;

        let mut child = Command::new("sh")
            .arg("-c")
            .arg(command)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .current_dir(&self.project_root)
            .spawn()?;

        if let Some(mut stdin) = child.stdin.take() {
            stdin.write_all(input_json.as_bytes()).await?;
            stdin.flush().await?;
            drop(stdin);
        }

        let result = if let Some(ms) = timeout_ms {
            timeout(Duration::from_millis(ms), child.wait_with_output()).await
        } else {
            Ok(child.wait_with_output().await)
        };

        match result {
            Ok(Ok(output)) => {
                let exit_code = output.status.code();
                let stdout = String::from_utf8_lossy(&output.stdout);

                // Try to parse JSON output from stdout
                if !stdout.trim().is_empty() {
                    if let Ok(hook_output) = serde_json::from_str::<HookOutput>(stdout.trim()) {
                        return Ok(hook_output);
                    }
                }

                // Fall back to exit code interpretation
                match exit_code {
                    Some(0) => Ok(HookOutput::allow()),
                    Some(2) => Ok(HookOutput::deny()),
                    Some(code) => Err(eyre!(
                        "Hook {} failed with exit code {}: {}",
                        event.hook_name(),
                        code,
                        String::from_utf8_lossy(&output.stderr)
                    )),
                    None => Err(eyre!("Hook {} terminated by signal", event.hook_name())),
                }
            }
            Ok(Err(e)) => Err(eyre!("Failed to execute hook {}: {}", event.hook_name(), e)),
            Err(_) => Err(eyre!("Hook {} timed out", event.hook_name())),
        }
    }

    /// Merge two hook outputs, with the second taking precedence
    fn merge_outputs(base: HookOutput, new: HookOutput) -> HookOutput {
        HookOutput {
            continue_execution: base.continue_execution && new.continue_execution,
            suppress_output: base.suppress_output || new.suppress_output,
            system_message: new.system_message.or(base.system_message),
            permission_decision: new.permission_decision.or(base.permission_decision),
            updated_input: new.updated_input.or(base.updated_input),
        }
    }

    fn get_hook_path(&self, event: &HookEvent) -> PathBuf {
        self.project_root
            .join(".owlen")
            .join("hooks")
            .join(event.hook_name())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn hook_event_serializes_correctly() {
        let event = HookEvent::PreToolUse {
            tool: "Read".to_string(),
            args: serde_json::json!({"path": "/tmp/test.txt"}),
        };

        let json = serde_json::to_string(&event).unwrap();
        assert!(json.contains("\"event\":\"preToolUse\""));
        assert!(json.contains("\"tool\":\"Read\""));
    }

    #[test]
    fn hook_event_names() {
        assert_eq!(
            HookEvent::PreToolUse {
                tool: "Read".to_string(),
                args: serde_json::json!({}),
            }
            .hook_name(),
            "PreToolUse"
        );
        assert_eq!(
            HookEvent::SessionStart {
                session_id: "123".to_string(),
            }
            .hook_name(),
            "SessionStart"
        );
        assert_eq!(
            HookEvent::Stop {
                reason: "task_complete".to_string(),
                num_messages: 10,
                num_tool_calls: 5,
            }
            .hook_name(),
            "Stop"
        );
        assert_eq!(
            HookEvent::SubagentStop {
                agent_id: "abc123".to_string(),
                agent_type: "explore".to_string(),
                reason: "completed".to_string(),
            }
            .hook_name(),
            "SubagentStop"
        );
    }

    #[test]
    fn stop_event_serializes_correctly() {
        let event = HookEvent::Stop {
            reason: "task_complete".to_string(),
            num_messages: 10,
            num_tool_calls: 5,
        };

        let json = serde_json::to_string(&event).unwrap();
        assert!(json.contains("\"event\":\"stop\""));
        assert!(json.contains("\"reason\":\"task_complete\""));
        assert!(json.contains("\"numMessages\":10"));
        assert!(json.contains("\"numToolCalls\":5"));
    }

    #[test]
    fn hook_output_defaults() {
        let output = HookOutput::default();
        assert!(output.continue_execution);
        assert!(!output.suppress_output);
        assert!(output.system_message.is_none());
        assert!(output.permission_decision.is_none());
    }

    #[test]
    fn hook_output_builders() {
        let output = HookOutput::allow()
            .with_system_message("Test message")
            .with_permission(HookPermission::Allow);

        assert!(output.continue_execution);
        assert_eq!(output.system_message, Some("Test message".to_string()));
        assert_eq!(output.permission_decision, Some(HookPermission::Allow));

        let deny = HookOutput::deny();
        assert!(!deny.continue_execution);
    }

    #[test]
    fn hook_output_deserializes() {
        let json = r#"{"continueExecution": true, "suppressOutput": false, "systemMessage": "Hello"}"#;
        let output: HookOutput = serde_json::from_str(json).unwrap();
        assert!(output.continue_execution);
        assert!(!output.suppress_output);
        assert_eq!(output.system_message, Some("Hello".to_string()));
    }

    #[test]
    fn hook_output_to_result() {
        assert_eq!(HookOutput::allow().to_result(), HookResult::Allow);
        assert_eq!(HookOutput::deny().to_result(), HookResult::Deny);
    }
}