use crate::dsl::schema::HookCommand;
use crate::error::{Error, Result};
use std::process::Stdio;
use tokio::process::Command;
#[derive(Debug, Clone)]
pub struct HookContext {
pub workflow_name: String,
pub stage: Option<String>,
pub error: Option<String>,
}
impl HookContext {
pub fn new(workflow_name: String) -> Self {
Self {
workflow_name,
stage: None,
error: None,
}
}
pub fn with_stage(mut self, stage: String) -> Self {
self.stage = Some(stage);
self
}
pub fn with_error(mut self, error: String) -> Self {
self.error = Some(error);
self
}
}
pub struct HooksExecutor;
impl HooksExecutor {
pub async fn execute_hooks(hooks: &[HookCommand], context: &HookContext) -> Result<()> {
for hook in hooks {
Self::execute_hook(hook, context).await?;
}
Ok(())
}
async fn execute_hook(hook: &HookCommand, context: &HookContext) -> Result<()> {
let (command, description) = match hook {
HookCommand::Command(cmd) => (cmd.as_str(), None),
HookCommand::CommandSpec {
command,
description,
} => (command.as_str(), description.as_ref()),
};
println!(
"Executing hook{}: {}",
description.map(|d| format!(" ({})", d)).unwrap_or_default(),
command
);
let output = Command::new("sh")
.arg("-c")
.arg(command)
.env("WORKFLOW_NAME", &context.workflow_name)
.env("WORKFLOW_STAGE", context.stage.as_deref().unwrap_or(""))
.env("WORKFLOW_ERROR", context.error.as_deref().unwrap_or(""))
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.output()
.await
.map_err(|e| Error::InvalidInput(format!("Failed to execute hook: {}", e)))?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
return Err(Error::InvalidInput(format!(
"Hook command failed: {}",
stderr
)));
}
let stdout = String::from_utf8_lossy(&output.stdout);
if !stdout.trim().is_empty() {
println!("Hook output: {}", stdout.trim());
}
Ok(())
}
pub async fn execute_pre_workflow(hooks: &[HookCommand], workflow_name: &str) -> Result<()> {
let context = HookContext::new(workflow_name.to_string());
println!("Running pre-workflow hooks...");
Self::execute_hooks(hooks, &context).await
}
pub async fn execute_post_workflow(hooks: &[HookCommand], workflow_name: &str) -> Result<()> {
let context = HookContext::new(workflow_name.to_string());
println!("Running post-workflow hooks...");
Self::execute_hooks(hooks, &context).await
}
pub async fn execute_stage_complete(
hooks: &[HookCommand],
workflow_name: &str,
stage: &str,
) -> Result<()> {
let context = HookContext::new(workflow_name.to_string()).with_stage(stage.to_string());
println!("Running stage complete hooks for '{}'...", stage);
Self::execute_hooks(hooks, &context).await
}
pub async fn execute_error(
hooks: &[HookCommand],
workflow_name: &str,
error: &str,
) -> Result<()> {
let context = HookContext::new(workflow_name.to_string()).with_error(error.to_string());
println!("Running error hooks...");
Self::execute_hooks(hooks, &context).await
}
}
#[derive(Debug, Clone)]
pub enum RecoveryStrategy {
Retry {
max_attempts: u32,
config: Option<crate::dsl::schema::ErrorHandlingSpec>,
},
Skip,
Fallback { agent_name: String },
Abort,
}
pub struct ErrorRecovery;
impl ErrorRecovery {
pub fn get_strategy(
retry_count: Option<u32>,
fallback_agent: Option<&str>,
) -> RecoveryStrategy {
if let Some(agent) = fallback_agent {
RecoveryStrategy::Fallback {
agent_name: agent.to_string(),
}
} else if let Some(max_attempts) = retry_count {
RecoveryStrategy::Retry {
max_attempts,
config: None, }
} else {
RecoveryStrategy::Abort
}
}
pub fn get_strategy_from_spec(
spec: Option<&crate::dsl::schema::ErrorHandlingSpec>,
) -> RecoveryStrategy {
if let Some(error_spec) = spec {
if let Some(agent) = &error_spec.fallback_agent {
RecoveryStrategy::Fallback {
agent_name: agent.clone(),
}
} else {
let has_retry = error_spec.retry > 0;
if has_retry {
RecoveryStrategy::Retry {
max_attempts: error_spec.retry,
config: Some(error_spec.clone()),
}
} else {
RecoveryStrategy::Abort
}
}
} else {
RecoveryStrategy::Abort
}
}
pub fn should_retry(strategy: &RecoveryStrategy, current_attempt: u32) -> bool {
match strategy {
RecoveryStrategy::Retry { max_attempts, .. } => current_attempt <= *max_attempts,
_ => false,
}
}
pub fn get_fallback_agent(strategy: &RecoveryStrategy) -> Option<&str> {
match strategy {
RecoveryStrategy::Fallback { agent_name } => Some(agent_name),
_ => None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_hook_context_creation() {
let context = HookContext::new("test_workflow".to_string());
assert_eq!(context.workflow_name, "test_workflow");
assert!(context.stage.is_none());
assert!(context.error.is_none());
}
#[test]
fn test_hook_context_with_stage() {
let context =
HookContext::new("test_workflow".to_string()).with_stage("stage1".to_string());
assert_eq!(context.stage, Some("stage1".to_string()));
}
#[test]
fn test_hook_context_with_error() {
let context =
HookContext::new("test_workflow".to_string()).with_error("error msg".to_string());
assert_eq!(context.error, Some("error msg".to_string()));
}
#[tokio::test]
async fn test_execute_simple_hook() {
let hooks = vec![HookCommand::Command("echo 'test'".to_string())];
let context = HookContext::new("test".to_string());
let result = HooksExecutor::execute_hooks(&hooks, &context).await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_execute_hook_with_description() {
let hooks = vec![HookCommand::CommandSpec {
command: "echo 'test'".to_string(),
description: Some("Test hook".to_string()),
}];
let context = HookContext::new("test".to_string());
let result = HooksExecutor::execute_hooks(&hooks, &context).await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_execute_failing_hook() {
let hooks = vec![HookCommand::Command("exit 1".to_string())];
let context = HookContext::new("test".to_string());
let result = HooksExecutor::execute_hooks(&hooks, &context).await;
assert!(result.is_err());
}
#[test]
fn test_recovery_strategy_retry() {
let strategy = ErrorRecovery::get_strategy(Some(3), None);
match strategy {
RecoveryStrategy::Retry { max_attempts, .. } => {
assert_eq!(max_attempts, 3)
}
_ => panic!("Wrong strategy"),
}
}
#[test]
fn test_recovery_strategy_fallback() {
let strategy = ErrorRecovery::get_strategy(None, Some("fallback_agent"));
match strategy {
RecoveryStrategy::Fallback { agent_name } => {
assert_eq!(agent_name, "fallback_agent")
}
_ => panic!("Wrong strategy"),
}
}
#[test]
fn test_should_retry() {
let strategy = RecoveryStrategy::Retry {
max_attempts: 3,
config: None,
};
assert!(ErrorRecovery::should_retry(&strategy, 0)); assert!(ErrorRecovery::should_retry(&strategy, 1)); assert!(ErrorRecovery::should_retry(&strategy, 2)); assert!(ErrorRecovery::should_retry(&strategy, 3)); assert!(!ErrorRecovery::should_retry(&strategy, 4)); }
}