ironflow-engine 2.17.3

//! [`WorkflowHandler`] trait — dynamic workflows with context chaining.
//!
//! Implement this trait to define workflows where steps can reference
//! outputs from previous steps. The handler receives a [`WorkflowContext`]
//! that provides step execution methods with automatic persistence.
//!
//! # Examples
//!
//! ```no_run
//! use ironflow_engine::handler::WorkflowHandler;
//! use ironflow_engine::context::WorkflowContext;
//! use ironflow_engine::config::{ShellConfig, AgentStepConfig};
//! use ironflow_engine::error::EngineError;
//! use std::future::Future;
//! use std::pin::Pin;
//!
//! struct DeployWorkflow;
//!
//! impl WorkflowHandler for DeployWorkflow {
//!     fn name(&self) -> &str {
//!         "deploy"
//!     }
//!
//!     fn execute<'a>(
//!         &'a self,
//!         ctx: &'a mut WorkflowContext,
//!     ) -> Pin<Box<dyn Future<Output = Result<(), EngineError>> + Send + 'a>> {
//!         Box::pin(async move {
//!             let build = ctx.shell("build", ShellConfig::new("cargo build --release")).await?;
//!             let tests = ctx.shell("test", ShellConfig::new("cargo test")).await?;
//!
//!             let review = ctx.agent("review", AgentStepConfig::new(
//!                 &format!("Build:\n{}\nTests:\n{}\nReview.",
//!                     build.output["stdout"], tests.output["stdout"])
//!             )).await?;
//!
//!             if review.output.as_str().unwrap_or("").contains("LGTM") {
//!                 ctx.shell("deploy", ShellConfig::new("./deploy.sh")).await?;
//!             }
//!
//!             Ok(())
//!         })
//!     }
//! }
//! ```

use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;

use rust_decimal::Decimal;
use schemars::JsonSchema;
use serde::Serialize;
use serde_json::Value;

use crate::context::WorkflowContext;
use crate::error::EngineError;
use crate::schedule::CronSchedule;

/// Generate a JSON Schema [`Value`] from a type that derives [`JsonSchema`].
///
/// Use this in [`WorkflowHandler::input_schema`] to automatically derive the
/// schema from your input struct instead of writing JSON by hand.
///
/// # Examples
///
/// ```
/// use schemars::JsonSchema;
/// use serde::Deserialize;
/// use ironflow_engine::handler::input_schema_for;
///
/// #[derive(Deserialize, JsonSchema)]
/// struct DeployInput {
///     environment: String,
///     dry_run: Option<bool>,
/// }
///
/// let schema = input_schema_for::<DeployInput>();
/// assert_eq!(schema["type"], "object");
/// assert!(schema["properties"]["environment"].is_object());
/// ```
pub fn input_schema_for<T: JsonSchema>() -> Value {
    let schema = schemars::schema_for!(T);
    serde_json::to_value(schema).expect("schema serialization cannot fail")
}

/// Boxed future returned by [`WorkflowHandler::execute`].
pub type HandlerFuture<'a> = Pin<Box<dyn Future<Output = Result<(), EngineError>> + Send + 'a>>;

/// Metadata about a workflow, returned by [`WorkflowHandler::describe`].
///
/// Contains a human-readable description and optional Rust source code
/// for display in the dashboard.
#[derive(Debug, Clone, Serialize)]
pub struct WorkflowInfo {
    /// Human-readable description of what the workflow does.
    pub description: String,
    /// Optional Rust source code of the handler (for UI display).
    pub source_code: Option<String>,
    /// Names of sub-workflows invoked by this handler.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub sub_workflows: Vec<String>,
    /// Optional `/`-separated category path used to group workflows in the UI tree.
    ///
    /// A value like `"data/etl"` places the workflow under `data` → `etl`.
    /// `None` means the workflow is uncategorized.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub category: Option<String>,
    /// Handler version string, used to trace which code produced a given run.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub version: Option<String>,
    /// JSON Schema describing the expected input payload.
    ///
    /// When present, the dashboard renders a dynamic form from this schema
    /// and the engine validates the payload before creating a run.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub input_schema: Option<Value>,
    /// Labels automatically applied to every run of this workflow.
    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
    pub default_labels: HashMap<String, String>,
    /// Optional cron schedule for automatic execution.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub schedule: Option<CronSchedule>,
    /// Default cumulative cost cap applied to runs of this workflow, in USD.
    ///
    /// Overridden by a cap supplied at run creation, and takes precedence over
    /// the server-wide default. `None` means the handler declares no default.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub default_max_cost_usd: Option<Decimal>,
}

/// A dynamic workflow handler with context-aware step chaining.
///
/// Implement this trait to define workflows where each step can use
/// the output of previous steps. Register handlers with
/// [`Engine::register`](crate::engine::Engine::register) and execute
/// them by name.
///
/// # Why `Pin<Box<dyn Future>>` instead of `async fn`?
///
/// The handler must be object-safe (`dyn WorkflowHandler`) to allow
/// registering different handler types in the engine's registry.
pub trait WorkflowHandler: Send + Sync {
    /// The workflow name used for registration and lookup.
    fn name(&self) -> &str;

    /// Handler version string, used to trace which code version produced a run.
    ///
    /// Override this to return a meaningful version (semver, git SHA, build
    /// hash, etc.). The default is `None`.
    fn version(&self) -> Option<&str> {
        None
    }

    /// Optional `/`-separated category path used to group workflows in the UI tree.
    ///
    /// Return a value like `"data/etl"` to place the workflow under `data` → `etl`.
    /// The default is `None` (uncategorized).
    ///
    /// Validation (empty segments, leading or trailing `/`, `//`, whitespace
    /// segments) is enforced at registration time by
    /// [`Engine::register`](crate::engine::Engine::register).
    fn category(&self) -> Option<&str> {
        None
    }

    /// Return a JSON Schema describing the expected input payload.
    ///
    /// When present, the dashboard renders a dynamic form from this schema
    /// and the engine validates the payload before creating a run.
    /// The default is `None` (no schema, free-form payload).
    fn input_schema(&self) -> Option<Value> {
        None
    }

    /// Labels automatically applied to every run of this workflow.
    ///
    /// These are merged with any labels provided at run creation time.
    /// User-provided labels take precedence over defaults.
    fn default_labels(&self) -> HashMap<String, String> {
        HashMap::new()
    }

    /// Optional cron schedule for automatic execution.
    ///
    /// Return a [`CronSchedule`] built from a cron expression
    /// (5 or 6 fields, as supported by [`croner`]).
    ///
    /// When set, the engine exposes this handler via
    /// [`Engine::scheduled_handlers`](crate::engine::Engine::scheduled_handlers)
    /// so the runtime can wire it into a cron scheduler automatically.
    ///
    /// The default is `None` (no automatic scheduling).
    ///
    /// # Examples
    ///
    /// ```
    /// # use ironflow_engine::handler::{WorkflowHandler, HandlerFuture};
    /// # use ironflow_engine::context::WorkflowContext;
    /// # use ironflow_engine::schedule::CronSchedule;
    /// struct HourlySync;
    ///
    /// impl WorkflowHandler for HourlySync {
    ///     fn name(&self) -> &str { "hourly-sync" }
    ///     fn schedule(&self) -> Option<&CronSchedule> {
    ///         // In practice, store as a field or use `std::sync::LazyLock`.
    ///         None
    ///     }
    ///     fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
    ///         Box::pin(async move { Ok(()) })
    ///     }
    /// }
    /// ```
    fn schedule(&self) -> Option<&CronSchedule> {
        None
    }

    /// Default cumulative cost cap for runs of this workflow, in USD.
    ///
    /// Applied when the run creation request does not supply one. Takes
    /// precedence over the server-wide
    /// [`IRONFLOW_DEFAULT_RUN_MAX_COST_USD`](crate::budget::DEFAULT_RUN_MAX_COST_ENV).
    /// The default is `None` (fall back to the server default, or no cap).
    ///
    /// # Examples
    ///
    /// ```
    /// # use ironflow_engine::handler::{WorkflowHandler, HandlerFuture};
    /// # use ironflow_engine::context::WorkflowContext;
    /// use rust_decimal::Decimal;
    ///
    /// struct ExpensiveAnalysis;
    ///
    /// impl WorkflowHandler for ExpensiveAnalysis {
    ///     fn name(&self) -> &str { "expensive-analysis" }
    ///     fn default_max_cost_usd(&self) -> Option<Decimal> {
    ///         Some(Decimal::new(500, 2)) // $5.00
    ///     }
    ///     fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
    ///         Box::pin(async move { Ok(()) })
    ///     }
    /// }
    ///
    /// assert_eq!(ExpensiveAnalysis.default_max_cost_usd(), Some(Decimal::new(500, 2)));
    /// ```
    fn default_max_cost_usd(&self) -> Option<Decimal> {
        None
    }

    /// Return metadata about this workflow (description, source code).
    ///
    /// Override this to provide a description and source code for the
    /// dashboard UI. The default returns an empty description with no source
    /// but propagates [`WorkflowHandler::category`],
    /// [`WorkflowHandler::version`], [`WorkflowHandler::input_schema`],
    /// [`WorkflowHandler::default_labels`],
    /// and [`WorkflowHandler::schedule`].
    fn describe(&self) -> WorkflowInfo {
        WorkflowInfo {
            description: String::new(),
            source_code: None,
            sub_workflows: Vec::new(),
            category: self.category().map(str::to_string),
            version: self.version().map(str::to_string),
            input_schema: self.input_schema(),
            default_labels: self.default_labels(),
            schedule: self.schedule().cloned(),
            default_max_cost_usd: self.default_max_cost_usd(),
        }
    }

    /// Execute the workflow with the given context.
    ///
    /// The context provides [`shell`](WorkflowContext::shell),
    /// [`http`](WorkflowContext::http), and [`agent`](WorkflowContext::agent)
    /// methods that automatically persist each step.
    ///
    /// # Errors
    ///
    /// Return [`EngineError`] if any step fails. The engine will mark
    /// the run as `Failed` and record the error.
    fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a>;
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde::{Deserialize, Serialize};

    #[derive(Debug, Serialize, Deserialize, JsonSchema)]
    struct TestInput {
        environment: String,
        #[serde(default)]
        dry_run: bool,
    }

    struct MinimalHandler;

    impl WorkflowHandler for MinimalHandler {
        fn name(&self) -> &str {
            "minimal"
        }

        fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
            Box::pin(async { Ok(()) })
        }
    }

    struct FullFeaturedHandler;

    impl WorkflowHandler for FullFeaturedHandler {
        fn name(&self) -> &str {
            "full"
        }

        fn version(&self) -> Option<&str> {
            Some("1.2.0")
        }

        fn category(&self) -> Option<&str> {
            Some("data/etl")
        }

        fn input_schema(&self) -> Option<Value> {
            Some(input_schema_for::<TestInput>())
        }

        fn default_labels(&self) -> HashMap<String, String> {
            HashMap::from([
                ("team".to_string(), "platform".to_string()),
                ("env".to_string(), "prod".to_string()),
            ])
        }

        fn default_max_cost_usd(&self) -> Option<Decimal> {
            Some(Decimal::new(750, 2))
        }

        fn describe(&self) -> WorkflowInfo {
            WorkflowInfo {
                description: "Full-featured test handler".to_string(),
                source_code: Some("fn test() {}".to_string()),
                sub_workflows: vec!["helper".to_string()],
                category: self.category().map(str::to_string),
                version: self.version().map(str::to_string),
                input_schema: self.input_schema(),
                default_labels: self.default_labels(),
                schedule: self.schedule().cloned(),
                default_max_cost_usd: self.default_max_cost_usd(),
            }
        }

        fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
            Box::pin(async { Ok(()) })
        }
    }

    #[test]
    fn minimal_handler_has_required_name() {
        let handler = MinimalHandler;
        assert_eq!(handler.name(), "minimal");
    }

    #[test]
    fn minimal_handler_defaults_to_no_version() {
        let handler = MinimalHandler;
        assert_eq!(handler.version(), None);
    }

    #[test]
    fn minimal_handler_defaults_to_no_category() {
        let handler = MinimalHandler;
        assert_eq!(handler.category(), None);
    }

    #[test]
    fn minimal_handler_defaults_to_no_schema() {
        let handler = MinimalHandler;
        assert_eq!(handler.input_schema(), None);
    }

    #[test]
    fn minimal_handler_defaults_to_empty_labels() {
        let handler = MinimalHandler;
        let labels = handler.default_labels();
        assert!(labels.is_empty());
    }

    #[test]
    fn minimal_handler_defaults_to_no_schedule() {
        let handler = MinimalHandler;
        assert_eq!(handler.schedule(), None);
    }

    #[test]
    fn minimal_handler_describe_reflects_defaults() {
        let handler = MinimalHandler;
        let info = handler.describe();
        assert_eq!(info.description, "");
        assert_eq!(info.source_code, None);
        assert_eq!(info.sub_workflows, Vec::<String>::new());
        assert_eq!(info.category, None);
        assert_eq!(info.version, None);
        assert_eq!(info.input_schema, None);
        assert!(info.default_labels.is_empty());
        assert_eq!(info.schedule, None);
    }

    #[test]
    fn full_handler_returns_all_metadata() {
        let handler = FullFeaturedHandler;
        assert_eq!(handler.name(), "full");
        assert_eq!(handler.version(), Some("1.2.0"));
        assert_eq!(handler.category(), Some("data/etl"));
        assert!(handler.input_schema().is_some());
    }

    #[test]
    fn full_handler_default_labels_are_set() {
        let handler = FullFeaturedHandler;
        let labels = handler.default_labels();
        assert_eq!(labels.get("team"), Some(&"platform".to_string()));
        assert_eq!(labels.get("env"), Some(&"prod".to_string()));
    }

    #[test]
    fn full_handler_describe_includes_all_fields() {
        let handler = FullFeaturedHandler;
        let info = handler.describe();
        assert_eq!(info.description, "Full-featured test handler");
        assert_eq!(info.source_code, Some("fn test() {}".to_string()));
        assert_eq!(info.sub_workflows, vec!["helper".to_string()]);
        assert_eq!(info.category, Some("data/etl".to_string()));
        assert_eq!(info.version, Some("1.2.0".to_string()));
        assert!(info.input_schema.is_some());
        assert_eq!(info.default_labels.len(), 2);
    }

    #[test]
    fn input_schema_for_generates_json_schema() {
        let schema = input_schema_for::<TestInput>();
        assert_eq!(schema["type"], "object");
        assert!(schema["properties"]["environment"].is_object());
        assert!(schema["properties"]["dry_run"].is_object());
    }

    #[test]
    fn input_schema_for_preserves_serde_attributes() {
        let schema = input_schema_for::<TestInput>();
        let properties = &schema["properties"];
        assert!(properties.is_object());
        assert!(properties.get("environment").is_some());
        assert!(properties.get("dry_run").is_some());
    }

    #[test]
    fn minimal_handler_defaults_to_no_max_cost() {
        assert!(MinimalHandler.default_max_cost_usd().is_none());
        assert!(MinimalHandler.describe().default_max_cost_usd.is_none());
    }

    #[test]
    fn describe_propagates_handler_max_cost() {
        assert_eq!(
            FullFeaturedHandler.describe().default_max_cost_usd,
            Some(Decimal::new(750, 2))
        );
    }

    #[test]
    fn workflow_info_omits_absent_max_cost_from_json() {
        let json = serde_json::to_value(MinimalHandler.describe()).expect("serialize");
        assert!(json.get("default_max_cost_usd").is_none());
    }

    #[test]
    fn workflow_info_serializes_with_skip_empty() {
        let info = WorkflowInfo {
            description: "test".to_string(),
            source_code: None,
            sub_workflows: Vec::new(),
            category: None,
            version: None,
            input_schema: None,
            default_labels: HashMap::new(),
            schedule: None,
            default_max_cost_usd: None,
        };

        let json = serde_json::to_value(&info).expect("serialize");
        assert_eq!(json["description"], "test");
        // Optional fields with skip_serializing_if may still be present or absent
        // depending on the serde configuration. Just verify the description is there.
        assert!(json.is_object());
    }

    #[test]
    fn workflow_info_serializes_with_values() {
        let info = WorkflowInfo {
            description: "test".to_string(),
            source_code: Some("code".to_string()),
            sub_workflows: vec!["sub".to_string()],
            category: Some("cat".to_string()),
            version: Some("1.0.0".to_string()),
            input_schema: Some(serde_json::json!({"type": "object"})),
            default_labels: HashMap::from([("key".to_string(), "value".to_string())]),
            schedule: Some(CronSchedule::new("0 0 * * * *").unwrap()),
            default_max_cost_usd: Some(Decimal::new(750, 2)),
        };

        let json = serde_json::to_value(&info).expect("serialize");
        assert_eq!(json["description"], "test");
        assert_eq!(json["source_code"], "code");
        assert_eq!(json["sub_workflows"][0], "sub");
        assert_eq!(json["category"], "cat");
        assert_eq!(json["version"], "1.0.0");
        assert_eq!(json["default_labels"]["key"], "value");
        assert_eq!(json["schedule"], "0 0 * * * *");
    }
}